add an internal API for obtaining current task id

2 files changed, 9 insertions, 0 deletions

diff --git a/modules/progress.py b/modules/progress.py
index c69ecf3d..05032ac5 100644
--- a/modules/progress.py
+++ b/modules/progress.py
@@ -4,6 +4,7 @@ import time
 

 import gradio as gr

 from pydantic import BaseModel, Field

+from typing import List

 

 from modules.shared import opts

 

@@ -37,6 +38,9 @@ def add_task_to_queue(id_job):
     pending_tasks[id_job] = time.time()

 

 

+class CurrentTaskResponse(BaseModel):

+  current_task: str = Field(default=None, title="Task ID", description="id of the current progress task")

+

 class ProgressRequest(BaseModel):

     id_task: str = Field(default=None, title="Task ID", description="id of the task to get progress for")

     id_live_preview: int = Field(default=-1, title="Live preview image ID", description="id of last received last preview image")

@@ -56,6 +60,8 @@ class ProgressResponse(BaseModel):
 def setup_progress_api(app):

     return app.add_api_route("/internal/progress", progressapi, methods=["POST"], response_model=ProgressResponse)

 

+def setup_current_task_api(app):

+    return app.add_api_route("/internal/current_task", current_task_api, methods=["GET"], response_model=CurrentTaskResponse)

 

 def progressapi(req: ProgressRequest):

     active = req.id_task == current_task

@@ -97,3 +103,5 @@ 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 current_task_api():

+  return CurrentTaskResponse(current_task=current_task)
\ No newline at end of file
diff --git a/webui.py b/webui.py
index b570895f..cca9c77f 100644
--- a/webui.py
+++ b/webui.py
@@ -279,6 +279,7 @@ def webui():
         setup_middleware(app)

 

         modules.progress.setup_progress_api(app)

+        modules.progress.setup_current_task_api(app)

 

         if launch_api:

             create_api(app)