diff --git a/modules/call_queue.py b/modules/call_queue.py index 92097c15..30ac26bc 100644 --- a/modules/call_queue.py +++ b/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() diff --git a/modules/progress.py b/modules/progress.py index 05032ac5..27a336ad 100644 --- a/modules/progress.py +++ b/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")