diff options
Diffstat (limited to 'modules/call_queue.py')
| -rw-r--r-- | modules/call_queue.py | 18 |
1 files changed, 13 insertions, 5 deletions
diff --git a/modules/call_queue.py b/modules/call_queue.py index 92097c15..43f6ebe0 100644 --- a/modules/call_queue.py +++ b/modules/call_queue.py @@ -4,10 +4,16 @@ import threading import traceback
import time
+import gradio as gr
from modules import shared, progress
queue_lock = threading.Lock()
+queue_lock_condition = threading.Condition(lock=queue_lock)
+def wrap_session_call(func):
+ def f(request: gr.Request, *args, **kwargs):
+ return func(request, *args, **kwargs)
+ return f
def wrap_queued_call(func):
def f(*args, **kwargs):
@@ -20,29 +26,31 @@ def wrap_queued_call(func): def wrap_gradio_gpu_call(func, extra_outputs=None):
- def f(*args, **kwargs):
+ def f(request: gr.Request, *args, **kwargs):
+ user = request.username
# if the first argument is a string that says "task(...)", it is treated as a job id
if len(args) > 0 and type(args[0]) == str and args[0][0:5] == "task(" and args[0][-1] == ")":
id_task = args[0]
- progress.add_task_to_queue(id_task)
+ progress.add_task_to_queue(user, id_task)
else:
id_task = None
with queue_lock:
shared.state.begin()
- progress.start_task(id_task)
+ progress.start_task(user, id_task)
try:
res = func(*args, **kwargs)
finally:
- progress.finish_task(id_task)
+ progress.finish_task(user, id_task)
+ progress.set_last_task_result(user, id_task, res)
shared.state.end()
return res
- return wrap_gradio_call(f, extra_outputs=extra_outputs, add_stats=True)
+ return wrap_session_call(wrap_gradio_call(f, extra_outputs=extra_outputs, add_stats=True))
def wrap_gradio_call(func, extra_outputs=None, add_stats=False):
|