main_thread.py

# This file is the main thread that handles all gradio calls for major t2i or i2i processing.
# Other gradio calls (like those from extensions) are not influenced.
# By using one single thread to process all major calls, model moving is significantly faster.


import threading
import time
import traceback

# Add global lock for model unloading
unload_lock = threading.Lock()
lock = threading.Lock()
last_id = 0
waiting_list = []
finished_list = []
last_exception = None
active_generations = 0  # Track active generation tasks
generation_lock = threading.Lock()

class Task:
    def __init__(self, task_id, func, args, kwargs):
        self.task_id = task_id
        self.func = func
        self.args = args
        self.kwargs = kwargs
        self.result = None
        self.exception = None

    def work(self):
        global last_exception, active_generations
        try:
            # Increment active generations counter
            with generation_lock:
                active_generations += 1
            
            self.result = self.func(*self.args, **self.kwargs)
            self.exception = None
            last_exception = None
            
        except Exception as e:
            traceback.print_exc()
            print(e)
            self.exception = e
            last_exception = e
        finally:
            # Decrement active generations counter
            with generation_lock:
                active_generations -= 1


def loop():
    global lock, last_id, waiting_list, finished_list
    while True:
        time.sleep(0.01)
        if len(waiting_list) > 0:
            with lock:
                task = waiting_list.pop(0)

            task.work()

            with lock:
                finished_list.append(task)

def wait_for_all_generations():
    """Wait until all generation tasks are complete"""
    while True:
        with generation_lock:
            if active_generations == 0:
                break
        time.sleep(0.1)

def async_run(func, *args, **kwargs):
    global lock, last_id, waiting_list, finished_list
    with lock:
        last_id += 1
        new_task = Task(task_id=last_id, func=func, args=args, kwargs=kwargs)
        waiting_list.append(new_task)
    return new_task.task_id


def run_and_wait_result(func, *args, **kwargs):
    global lock, last_id, waiting_list, finished_list
    current_id = async_run(func, *args, **kwargs)
    while True:
        time.sleep(0.01)
        finished_task = None
        for t in finished_list.copy():
            if t.task_id == current_id:
                finished_task = t
                break
        if finished_task is not None:
            with lock:
                finished_list.remove(finished_task)
            
            # Wait for all generations to complete before returning
            wait_for_all_generations()
            
            return finished_task.result