Merge pull request #7595 from siutin/feature/restore-progress

restore the progress from session lost / tab reload

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):