Browse Source

store the last generated result

siutin 2 years ago
parent
commit
9407f1731a
2 changed files with 11 additions and 0 deletions
  1. 1 0
      modules/call_queue.py
  2. 10 0
      modules/progress.py

+ 1 - 0
modules/call_queue.py

@@ -37,6 +37,7 @@ def wrap_gradio_gpu_call(func, extra_outputs=None):
                 res = func(*args, **kwargs)
             finally:
                 progress.finish_task(id_task)
+                progress.set_last_task_result(id_task, res)
 
             shared.state.end()
 

+ 10 - 0
modules/progress.py

@@ -37,6 +37,16 @@ def finish_task(id_task):
 def add_task_to_queue(id_job):
     pending_tasks[id_job] = time.time()
 
+last_task_id = None
+last_task_result = None
+
+def set_last_task_result(id_job, result):
+  global last_task_id
+  global last_task_result
+
+  last_task_id = id_job
+  last_task_result = result
+
 
 class CurrentTaskResponse(BaseModel):
   current_task: str = Field(default=None, title="Task ID", description="id of the current progress task")