scripts/comfyui.py

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