78 lines
2.1 KiB
Python
Executable File
78 lines
2.1 KiB
Python
Executable File
# 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 time
|
|
import traceback
|
|
import threading
|
|
|
|
|
|
lock = threading.Lock()
|
|
last_id = 0
|
|
waiting_list = []
|
|
finished_list = []
|
|
last_exception = None
|
|
|
|
|
|
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
|
|
try:
|
|
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
|
|
|
|
|
|
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 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(): # thread safe shallow copy without needing a lock
|
|
if t.task_id == current_id:
|
|
finished_task = t
|
|
break
|
|
if finished_task is not None:
|
|
with lock:
|
|
finished_list.remove(finished_task)
|
|
return finished_task.result
|
|
|