From bc3a4b4f40bd7eb9f111958fabd36b93bd495bc6 Mon Sep 17 00:00:00 2001 From: Konie Date: Wed, 25 Oct 2023 11:31:55 +0800 Subject: [PATCH 1/2] Add args for control private log and working queue --- fooocusapi/api.py | 2 +- fooocusapi/task_queue.py | 4 ++-- main.py | 28 +++++++++++++++++++--------- 3 files changed, 22 insertions(+), 12 deletions(-) diff --git a/fooocusapi/api.py b/fooocusapi/api.py index 586e980..569d3f6 100644 --- a/fooocusapi/api.py +++ b/fooocusapi/api.py @@ -11,7 +11,7 @@ app = FastAPI() -work_executor = ThreadPoolExecutor(max_workers=12, thread_name_prefix="worker_") +work_executor = ThreadPoolExecutor(max_workers=task_queue.queue_size*2, thread_name_prefix="worker_") img_generate_responses = { "200": { diff --git a/fooocusapi/task_queue.py b/fooocusapi/task_queue.py index f21130e..3002af2 100644 --- a/fooocusapi/task_queue.py +++ b/fooocusapi/task_queue.py @@ -48,7 +48,7 @@ class TaskQueue(object): def __init__(self, queue_size: int, hisotry_size: int): self.queue_size = queue_size - self.hisotry_size = hisotry_size + self.history_size = hisotry_size def add_task(self, type: TaskType, req_param: dict) -> QueueTask | None: """ @@ -99,6 +99,6 @@ def finish_task(self, seq: int): self.history.append(task) # Clean history - if len(self.history) > self.hisotry_size: + if len(self.history) > self.history_size: removed_task = self.history.pop(0) print(f"Clean task history, remove task: {removed_task.seq}") diff --git a/main.py b/main.py index b9fe9ba..da00a5d 100644 --- a/main.py +++ b/main.py @@ -246,6 +246,13 @@ def prepare_environments(args) -> bool: f"Invalid value for argument '--sync-repo', acceptable value are 'skip' and 'only'") exit(1) + import fooocusapi.worker as worker + worker.task_queue.queue_size = args.queue_size + worker.task_queue.history_size = args.queue_history + + if args.disable_private_log: + worker.save_log = False + if not skip_sync_repo: download_repositories() @@ -269,23 +276,23 @@ def prepare_environments(args) -> bool: return True -def pre_setup(skip_sync_repo: bool=False, disable_private_log: bool=False, load_all_models: bool=False, preload_pipeline: bool=False): +def pre_setup(skip_sync_repo: bool=False, disable_private_log: bool=False, load_all_models: bool=False, preload_pipeline: bool=False, skip_save_log: bool=False): class Args(object): sync_repo = None + disable_private_log = False preload_pipeline = False + queue_size = 3 + queue_history = 6 print("[Pre Setup] Prepare environments") args = Args() + args.disable_private_log = disable_private_log args.preload_pipeline = preload_pipeline if skip_sync_repo: args.sync_repo = 'skip' prepare_environments(args) - if disable_private_log: - import fooocusapi.worker as worker - worker.save_log = False - if load_all_models: import modules.path as path from fooocusapi.parameters import inpaint_model_version @@ -309,14 +316,17 @@ def ini_cbh_args(): parser = argparse.ArgumentParser() parser.add_argument("--port", type=int, default=8888, - help="Set the listen port") + help="Set the listen port, default: 8888") parser.add_argument("--host", type=str, - default='127.0.0.1', help="Set the listen host") + default='127.0.0.1', help="Set the listen host, default: 127.0.0.1") parser.add_argument("--log-level", type=str, - default='info', help="Log info for Uvicorn") + default='info', help="Log info for Uvicorn, default: info") parser.add_argument("--sync-repo", default=None, help="Sync dependent git repositories to local, 'skip' for skip sync action, 'only' for only do the sync action and not launch app") - parser.add_argument("--preload-pipeline", default=False, action="store_true", help="True for preload pipeline before start http server") + parser.add_argument("--disable-private-log", default=False, action="store_true", help="Disable Fooocus private log, won't save output files (include generated image files)") + parser.add_argument("--preload-pipeline", default=False, action="store_true", help="Preload pipeline before start http server") + parser.add_argument("--queue-size", type=int, default=3, help="Working queue size, default: 3, generation requests exceeding working queue size will return failure") + parser.add_argument("--queue-history", type=int, default=6, help="Finished jobs reserve in memory size, default: 6") args = parser.parse_args() From 376800ed499bca8e0df2f5624d005bc538875175 Mon Sep 17 00:00:00 2001 From: Konie Date: Wed, 25 Oct 2023 11:42:31 +0800 Subject: [PATCH 2/2] Add `Query Job Queue Info` api --- README.md | 7 ++++++- fooocus_api_version.py | 2 +- fooocusapi/api.py | 7 ++++++- fooocusapi/models.py | 8 +++++++- 4 files changed, 20 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 81aab68..468ced8 100644 --- a/README.md +++ b/README.md @@ -94,4 +94,9 @@ Alternative api for 'Image Prompt' tab of Fooocus Gradio interface. #### Query Job > GET /v1/generation/query-job -Query async generation request results, return job progress and generation results. \ No newline at end of file +Query async generation request results, return job progress and generation results. + +#### Query Job Queue Info +> GET /v1/generation/job-queue + +Query job queue info, include running job count, finished job count and last job id. \ No newline at end of file diff --git a/fooocus_api_version.py b/fooocus_api_version.py index 2f8d1fc..9bc95e8 100644 --- a/fooocus_api_version.py +++ b/fooocus_api_version.py @@ -1 +1 @@ -version = '0.2.1' \ No newline at end of file +version = '0.2.2' \ No newline at end of file diff --git a/fooocusapi/api.py b/fooocusapi/api.py index 569d3f6..851634d 100644 --- a/fooocusapi/api.py +++ b/fooocusapi/api.py @@ -3,7 +3,7 @@ from fastapi.params import File import uvicorn from fooocusapi.api_utils import generation_output, req_to_params -from fooocusapi.models import AsyncJobResponse, GeneratedImageBase64, ImgInpaintOrOutpaintRequest, ImgPromptRequest, ImgUpscaleOrVaryRequest, Text2ImgRequest +from fooocusapi.models import AsyncJobResponse, GeneratedImageBase64, 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 @@ -147,6 +147,11 @@ def query_job(job_id: int): return generation_output(queue_task, False) +@app.get("/v1/generation/job-queue", response_model=JobQueueInfo, description="Query job queue info") +def job_queue(): + return JobQueueInfo(running_size=len(task_queue.queue), finished_size=len(task_queue.history), last_job_id=task_queue.last_seq) + + def start_app(args): uvicorn.run("fooocusapi.api:app", host=args.host, port=args.port, log_level=args.log_level) diff --git a/fooocusapi/models.py b/fooocusapi/models.py index cd3f74a..d432e68 100644 --- a/fooocusapi/models.py +++ b/fooocusapi/models.py @@ -361,4 +361,10 @@ class AsyncJobResponse(BaseModel): job_stage: AsyncJobStage job_progess: int job_status: str | None - job_result: List[GeneratedImageBase64] | None \ No newline at end of file + job_result: List[GeneratedImageBase64] | None + + +class JobQueueInfo(BaseModel): + running_size: int = Field(description="The current running and waiting job count") + finished_size: int = Field(description="Finished job cound (after auto clean)") + last_job_id: int = Field(description="Last submit generation job id") \ No newline at end of file