Add require_step_preivew param to Query Job api

mrhan1993 · Nov 15, 2023 · 519e19f · 519e19f
1 parent 82454f8
commit 519e19f
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 10 deletions.
diff --git a/docs/openapi.json b/docs/openapi.json
@@ -415,6 +415,16 @@
                             "type": "integer",
                             "title": "Job Id"
                         }
+                    },
+                    {
+                        "name": "require_step_preivew",
+                        "in": "query",
+                        "required": false,
+                        "schema": {
+                            "type": "boolean",
+                            "default": false,
+                            "title": "Require Step Preivew"
+                        }
                     }
                 ],
                 "responses": {
@@ -795,7 +805,7 @@
                             }
                         ],
                         "title": "Job Status",
-                        "description": "Job status in text"
+                        "description": "Job running status in text"
                     },
                     "job_step_preview": {
                         "anyOf": [
@@ -807,7 +817,7 @@
                             }
                         ],
                         "title": "Job Step Preview",
-                        "description": "Preview result in each generation step, as base64 image"
+                        "description": "Preview image of generation steps at current time, as base64 image"
                     },
                     "job_result": {
                         "anyOf": [

diff --git a/fooocus_api_version.py b/fooocus_api_version.py
@@ -1 +1 @@
-version = '0.3.13'
+version = '0.3.14'
diff --git a/fooocusapi/api.py b/fooocusapi/api.py
@@ -5,7 +5,7 @@
 import uvicorn
 from fooocusapi.api_utils import generation_output, req_to_params
 import fooocusapi.file_utils as file_utils
-from fooocusapi.models import AllModelNamesResponse, AsyncJobResponse,StopResponse , GeneratedImageResult, ImgInpaintOrOutpaintRequest, ImgPromptRequest, ImgUpscaleOrVaryRequest, JobQueueInfo, Text2ImgRequest
+from fooocusapi.models import AllModelNamesResponse, AsyncJobResponse, QueryJobRequest,StopResponse , GeneratedImageResult, ImgInpaintOrOutpaintRequest, ImgPromptRequest, ImgUpscaleOrVaryRequest, JobQueueInfo, Text2ImgRequest
 from fooocusapi.parameters import GenerationFinishReason, ImageGenerationResult
 from fooocusapi.task_queue import TaskType
 from fooocusapi.worker import process_generate, task_queue, process_top
@@ -147,12 +147,12 @@ def img_prompt(cn_img1: Optional[UploadFile] = File(None),
 
 
 @app.get("/v1/generation/query-job", response_model=AsyncJobResponse, description="Query async generation job")
-def query_job(job_id: int):
-    queue_task = task_queue.get_task(job_id, True)
+def query_job(req: QueryJobRequest=Depends()):
+    queue_task = task_queue.get_task(req.job_id, True)
     if queue_task is None:
         return Response(content="Job not found", status_code=404)
 
-    return generation_output(queue_task, False, False)
+    return generation_output(queue_task, streaming_output=False, require_base64=False, require_step_preivew=req.require_step_preivew)
 
 
 @app.get("/v1/generation/job-queue", response_model=JobQueueInfo, description="Query job queue info")

diff --git a/fooocusapi/api_utils.py b/fooocusapi/api_utils.py
@@ -125,7 +125,7 @@ def req_to_params(req: Text2ImgRequest) -> ImageGenerationParams:
                                  )
 
 
-def generation_output(results: QueueTask | List[ImageGenerationResult], streaming_output: bool, require_base64: bool) -> Response | List[GeneratedImageResult] | AsyncJobResponse:
+def generation_output(results: QueueTask | List[ImageGenerationResult], streaming_output: bool, require_base64: bool, require_step_preivew: bool=False) -> Response | List[GeneratedImageResult] | AsyncJobResponse:
     if isinstance(results, QueueTask):
         task = results
         job_stage = AsyncJobStage.running
@@ -143,12 +143,13 @@ def generation_output(results: QueueTask | List[ImageGenerationResult], streamin
                         task_result_require_base64 = True
 
                     job_result = generation_output(task.task_result, False, task_result_require_base64)
+        job_step_preview = None if not require_step_preivew else task.task_step_preview
         return AsyncJobResponse(job_id=task.seq,
                                 job_type=task.type,
                                 job_stage=job_stage,
                                 job_progess=task.finish_progess,
                                 job_status=task.task_status,
-                                job_step_preview=task.task_step_preview,
+                                job_step_preview=job_step_preview,
                                 job_result=job_result)
 
     if streaming_output:

diff --git a/fooocusapi/models.py b/fooocusapi/models.py
@@ -383,13 +383,18 @@ class AsyncJobStage(str, Enum):
     error = 'ERROR'
 
 
+class QueryJobRequest(BaseModel):
+    job_id: int = Field(description="Job ID to query")
+    require_step_preivew: bool = Field(False, description="Set to true will return preview image of generation steps at current time")
+
+
 class AsyncJobResponse(BaseModel):
     job_id: int = Field(description="Job ID")
     job_type: TaskType = Field(description="Job type")
     job_stage: AsyncJobStage = Field(description="Job running stage")
     job_progess: int = Field(description="Job running progress, 100 is for finished.")
     job_status: str | None = Field(None, description="Job running status in text")
-    job_step_preview: str | None = Field(None, description="Preview result in each generation step, as base64 image")
+    job_step_preview: str | None = Field(None, description="Preview image of generation steps at current time, as base64 image")
     job_result: List[GeneratedImageResult] | None = Field(None, description="Job generation result")