big rework of progressbar/preview system to allow multiple users to prompts at the same time and do not get previews of each other

1 files changed, 15 insertions, 4 deletions

diff --git a/modules/call_queue.py b/modules/call_queue.py
index 4cd49533..92097c15 100644
--- a/modules/call_queue.py
+++ b/modules/call_queue.py
@@ -4,7 +4,7 @@ import threading
 import traceback

 import time

 

-from modules import shared

+from modules import shared, progress

 

 queue_lock = threading.Lock()

 

@@ -22,12 +22,23 @@ def wrap_queued_call(func):
 def wrap_gradio_gpu_call(func, extra_outputs=None):

     def f(*args, **kwargs):

 

-        shared.state.begin()

+        # 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)

+        else:

+            id_task = None

 

         with queue_lock:

-            res = func(*args, **kwargs)

+            shared.state.begin()

+            progress.start_task(id_task)

+

+            try:

+                res = func(*args, **kwargs)

+            finally:

+                progress.finish_task(id_task)

 

-        shared.state.end()

+            shared.state.end()

 

         return res