Merge branch 'main' of https://github.com/konieshadow/Fooocus-API int…

…o main

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.
+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.

fooocus_api_version.py

@@ -1 +1 @@
-version = '0.2.1'
+version = '0.2.2'

fooocusapi/api.py

@@ -3,15 +3,15 @@
 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
 from concurrent.futures import ThreadPoolExecutor
 
 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": {
@@ -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)

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
+    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")

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}")

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()