fluxdev/stable-diffusion-webui-forge

1# This file is the main thread that handles all gradio calls for major t2i or i2i processing.2# Other gradio calls (like those from extensions) are not influenced.3# By using one single thread to process all major calls, model moving is significantly faster.4 5 6import time7import traceback8import threading9 10 11lock = threading.Lock()12last_id = 013waiting_list = []14finished_list = []15 16 17class Task:18    def __init__(self, task_id, func, args, kwargs):19        self.task_id = task_id20        self.func = func21        self.args = args22        self.kwargs = kwargs23        self.result = None24 25    def work(self):26        self.result = self.func(*self.args, **self.kwargs)27 28 29def loop():30    global lock, last_id, waiting_list, finished_list31    while True:32        time.sleep(0.01)33        if len(waiting_list) > 0:34            with lock:35                task = waiting_list.pop(0)36            try:37                task.work()38            except Exception as e:39                traceback.print_exc()40                print(e)41            with lock:42                finished_list.append(task)43 44 45def async_run(func, *args, **kwargs):46    global lock, last_id, waiting_list, finished_list47    with lock:48        last_id += 149        new_task = Task(task_id=last_id, func=func, args=args, kwargs=kwargs)50        waiting_list.append(new_task)51    return new_task.task_id52 53 54def run_and_wait_result(func, *args, **kwargs):55    global lock, last_id, waiting_list, finished_list56    current_id = async_run(func, *args, **kwargs)57    while True:58        time.sleep(0.01)59        finished_task = None60        for t in finished_list.copy():  # thread safe shallow copy without needing a lock61            if t.task_id == current_id:62                finished_task = t63                break64        if finished_task is not None:65            with lock:66                finished_list.remove(finished_task)67            return finished_task.result68 69