From 519e19fb69a7861bd25beb9453b33c249aeb0956 Mon Sep 17 00:00:00 2001 From: Konie Date: Wed, 15 Nov 2023 14:57:28 +0800 Subject: [PATCH] Add `require_step_preivew` param to `Query Job` api --- docs/openapi.json | 14 ++++++++++++-- fooocus_api_version.py | 2 +- fooocusapi/api.py | 8 ++++---- fooocusapi/api_utils.py | 5 +++-- fooocusapi/models.py | 7 ++++++- 5 files changed, 26 insertions(+), 10 deletions(-) diff --git a/docs/openapi.json b/docs/openapi.json index 7e38824..1eaef40 100644 --- 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 index cbd97d5..dc086e3 100644 --- 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 index 047b409..9321be4 100644 --- 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 index e05c236..254e7fe 100644 --- 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 index 3a68aa2..052b300 100644 --- 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")