alternate restore progress button implementation

1 files changed, 18 insertions, 0 deletions

diff --git a/modules/progress.py b/modules/progress.py
index c69ecf3d..5655346b 100644
--- a/modules/progress.py
+++ b/modules/progress.py
@@ -13,6 +13,8 @@ import modules.shared as shared
 current_task = None

 pending_tasks = {}

 finished_tasks = []

+recorded_results = []

+recorded_results_limit = 2

 

 

 def start_task(id_task):

@@ -33,6 +35,12 @@ def finish_task(id_task):
         finished_tasks.pop(0)

 

 

+def record_results(id_task, res):

+    recorded_results.append((id_task, res))

+    if len(recorded_results) > recorded_results_limit:

+        recorded_results.pop(0)

+

+

 def add_task_to_queue(id_job):

     pending_tasks[id_job] = time.time()

 

@@ -97,3 +105,13 @@ def progressapi(req: ProgressRequest):
 

     return ProgressResponse(active=active, queued=queued, completed=completed, progress=progress, eta=eta, live_preview=live_preview, id_live_preview=id_live_preview, textinfo=shared.state.textinfo)

 

+

+def restore_progress(id_task):

+    while id_task == current_task or id_task in pending_tasks:

+        time.sleep(0.1)

+

+    res = next(iter([x[1] for x in recorded_results if id_task == x[0]]), None)

+    if res is not None:

+        return res

+

+    return gr.update(), gr.update(), gr.update(), f"Couldn't restore progress for {id_task}: results either have been discarded or never were obtained"