File size: 5,615 Bytes
35cdf61 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 |
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' #test remove
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 |