From d8b90ac121cbf0c18b1dc9d56a5e1d14ca51e74e Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sun, 15 Jan 2023 18:50:56 +0300
Subject: big rework of progressbar/preview system to allow multiple users to
 prompts at the same time and do not get previews of each other

---
 modules/call_queue.py | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

(limited to 'modules/call_queue.py')

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
 
-- 
cgit v1.2.3