| import socket |
| import time |
| import requests |
| import subprocess |
| import os |
| import websocket |
| import uuid |
| import sqlite3 |
|
|
| def is_port_in_use(port, host='127.0.0.1'): |
| with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: |
| s.settimeout(1) |
| try: |
| s.bind((host, port)) |
| except socket.error: |
| return True |
| return False |
| |
| def setup_database(number_of_ports, starting_port=50000): |
| conn = sqlite3.connect('./scripts/ports.db') |
| cursor = conn.cursor() |
| cursor.execute('DROP TABLE IF EXISTS port_to_use') |
| cursor.execute('DROP TABLE IF EXISTS port_in_use') |
| cursor.execute(''' |
| CREATE TABLE port_to_use ( |
| port INTEGER |
| ) |
| ''') |
| cursor.execute(''' |
| CREATE TABLE port_in_use ( |
| port INTEGER |
| ) |
| ''') |
|
|
| count = 0 |
| con = 0 |
|
|
| while(True): |
|
|
| if not is_port_in_use(starting_port + count) : |
| cursor.execute('INSERT INTO port_to_use (port) VALUES (?)', (starting_port + count,)) |
| con = con + 1 |
| |
| count = count + 1 |
| if con >= number_of_ports: |
| break |
| |
| conn.commit() |
| conn.close() |
| |
| class ComfyUI: |
| def __init__(self): |
| |
| self.port = None |
| self.process = None |
|
|
| def get_ports_to_use(self): |
| conn = sqlite3.connect('./scripts/ports.db') |
| cursor = conn.cursor() |
| cursor.execute('SELECT port FROM port_to_use ORDER BY port ASC LIMIT 1') |
| first_port = cursor.fetchone() |
| if not first_port: |
| return None |
| first_port = first_port[0] |
| cursor.execute('DELETE FROM port_to_use WHERE port = ?', (first_port,)) |
| conn.commit() |
| conn.close() |
| return first_port |
| |
| def search_port_in_use(self, port): |
| conn = sqlite3.connect('./scripts/ports.db') |
| cursor = conn.cursor() |
| cursor.execute('SELECT * FROM port_in_use WHERE port = ?', (port,)) |
| rows = cursor.fetchall() |
| conn.close() |
| return len(rows) > 0 |
| |
| def insert_port_in_use(self, port): |
| conn = sqlite3.connect('./scripts/ports.db') |
| cursor = conn.cursor() |
| cursor.execute('INSERT INTO port_in_use (port) VALUES (?)', (port,)) |
| conn.commit() |
| conn.close() |
|
|
| def insert_port_to_use(self, port): |
| conn = sqlite3.connect('./scripts/ports.db') |
| cursor = conn.cursor() |
| cursor.execute('INSERT INTO port_to_use (port) VALUES (?)', (port,)) |
| conn.commit() |
| conn.close() |
|
|
| def delete_port_in_use(self, port): |
| conn = sqlite3.connect('./scripts/ports.db') |
| cursor = conn.cursor() |
| cursor.execute('DELETE FROM port_in_use WHERE port = ?', (port,)) |
| conn.commit() |
| conn.close() |
|
|
| def is_server_ready(self, url, timeout=30): |
| start_time = time.time() |
| client_id = str(uuid.uuid4()) |
| |
| while True: |
| ws = None |
| try: |
| ws = websocket.WebSocket() |
| ws.connect("ws://{}/ws?clientId={}".format(url, client_id)) |
| try: |
| ws.send('ping') |
| response = ws.recv() |
| if response: |
| ws.close() |
| return True |
| except websocket.WebSocketException: |
| pass |
| except (websocket.WebSocketException, ConnectionRefusedError) as e: |
| time.sleep(1) |
| if time.time() - start_time > timeout: |
| return False |
| time.sleep(1) |
|
|
| def UI(self): |
| port_to_use = self.get_ports_to_use() |
| if port_to_use is None: |
| port_to_check = 60000 |
| con_count = 0 |
| while(True): |
| if con_count > 500: |
| return False |
| if is_port_in_use(port_to_check) or self.search_port_in_use(port_to_check): |
| port_to_check = port_to_check + 1 |
| con_count = con_count + 1 |
| else: |
| break |
| self.insert_port_in_use(port_to_check) |
| port_to_use = port_to_check |
| |
| script_dir = "./ComfyUI" |
| script_dir = os.path.abspath(script_dir) |
| |
| code_dir = "./ComfyUI/main.py" |
| code_dir = os.path.abspath(code_dir) |
| |
| env_path = os.getcwd() + '/AI_ENV/bin/python' |
| env_path = '/home/itek/miniconda3/envs/ai/bin/python' |
| command = [env_path, os.path.abspath(code_dir), "--port", str(port_to_use), '--listen' , '127.0.0.1'] |
| self.process = subprocess.Popen(command, cwd=script_dir, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) |
| port = str(port_to_use) |
| if self.is_server_ready(f"localhost:{port}", timeout=20): |
| self.port = port_to_use |
| self.insert_port_in_use(self.port) |
| return True |
| return False |
|
|
| def restart_if_down(self): |
|
|
| if self.port is None: |
| return False |
|
|
| if not is_port_in_use(self.port): |
| |
| if self.process is not None: |
| self.process.terminate() |
| self.process.kill() |
|
|
| self.insert_port_to_use(self.port) |
| self.delete_port_in_use(self.port) |
| self.process = None |
| self.port = None |
|
|
| return self.UI() |
|
|
| def end(self): |
| self.delete_port_in_use(self.port) |
| self.insert_port_to_use(self.port) |
| self.port = None |
|
|
| self.process.terminate() |
| self.process.kill() |
| self.process = None |
