Update fooocusapi/routes/generate_v1.py

fooocusapi/routes/generate_v1.py

@@ -1,5 +1,4 @@
 """Generate API V1 routes
-
 """
 from typing import List, Optional
 from fastapi import APIRouter, Depends, Header, Query, UploadFile
@@ -39,10 +38,10 @@ def stop_worker():
 
 
 @secure_router.post(
-        path="/v1/generation/text-to-image",
-        response_model=List[GeneratedImageResult] | AsyncJobResponse,
-        responses=img_generate_responses,
-        tags=["GenerateV1"])
+    path="/v1/generation/text-to-image",
+    response_model=List[GeneratedImageResult] | AsyncJobResponse,
+    responses=img_generate_responses,
+    tags=["GenerateV1"])
 def text2img_generation(
     req: Text2ImgRequest,
     accept: str = Header(None),
@@ -65,10 +64,10 @@ def text2img_generation(
 
 
 @secure_router.post(
-        path="/v1/generation/image-upscale-vary",
-        response_model=List[GeneratedImageResult] | AsyncJobResponse,
-        responses=img_generate_responses,
-        tags=["GenerateV1"])
+    path="/v1/generation/image-upscale-vary",
+    response_model=List[GeneratedImageResult] | AsyncJobResponse,
+    responses=img_generate_responses,
+    tags=["GenerateV1"])
 def img_upscale_or_vary(
     input_image: UploadFile,
     req: ImgUpscaleOrVaryRequest = Depends(ImgUpscaleOrVaryRequest.as_form),
@@ -93,10 +92,10 @@ def img_upscale_or_vary(
 
 
 @secure_router.post(
-        path="/v1/generation/image-inpaint-outpaint",
-        response_model=List[GeneratedImageResult] | AsyncJobResponse,
-        responses=img_generate_responses,
-        tags=["GenerateV1"])
+    path="/v1/generation/image-inpaint-outpaint",
+    response_model=List[GeneratedImageResult] | AsyncJobResponse,
+    responses=img_generate_responses,
+    tags=["GenerateV1"])
 def img_inpaint_or_outpaint(
     input_image: UploadFile,
     req: ImgInpaintOrOutpaintRequest = Depends(ImgInpaintOrOutpaintRequest.as_form),
@@ -119,12 +118,15 @@ def img_inpaint_or_outpaint(
 
 
 @secure_router.post(
-        path="/v1/generation/image-prompt",
-        response_model=List[GeneratedImageResult] | AsyncJobResponse,
-        responses=img_generate_responses,
-        tags=["GenerateV1"])
+    path="/v1/generation/image-prompt",
+    response_model=List[GeneratedImageResult] | AsyncJobResponse,
+    responses=img_generate_responses,
+    tags=["GenerateV1"])
 def img_prompt(
-    cn_img1: Optional[UploadFile] = File(None),
+    cn_img1: UploadFile,  # Required field
+    cn_img2: Optional[UploadFile] = File(None),
+    cn_img3: Optional[UploadFile] = File(None),
+    cn_img4: Optional[UploadFile] = File(None),
     req: ImgPromptRequest = Depends(ImgPromptRequest.as_form),
     accept: str = Header(None),
     accept_query: str | None = Query(
@@ -132,10 +134,13 @@ def img_prompt(
         description="Parameter to override 'Accept' header, 'image/png' for output bytes")):
     """\nImage Prompt\n
     Image Prompt
-    A prompt-based image generation.
+    A prompt-based image generation with up to four control images.
     Arguments:
-        cn_img1 {UploadFile} -- Input image file
-        req {ImgPromptRequest} -- Request body
+        cn_img1 {UploadFile} -- First control image (required)
+        cn_img2 {UploadFile} -- Second control image (optional)
+        cn_img3 {UploadFile} -- Third control image (optional)
+        cn_img4 {UploadFile} -- Fourth control image (optional)
+        req {ImgPromptRequest} -- Request body containing prompt and control parameters
         accept {str} -- Accept header
         accept_query {str} -- Parameter to override 'Accept' header, 'image/png' for output bytes
     Returns:
@@ -148,22 +153,21 @@ def img_prompt(
 
 
 @secure_router.post(
-        path="/v1/generation/image-enhance",
-        response_model=List[GeneratedImageResult] | AsyncJobResponse,
-        responses=img_generate_responses,
-        tags=["GenerateV1"])
-def img_prompt(
+    path="/v1/generation/image-enhance",
+    response_model=List[GeneratedImageResult] | AsyncJobResponse,
+    responses=img_generate_responses,
+    tags=["GenerateV1"])
+def img_enhance(  # Renamed to avoid function name collision
     enhance_input_image: Optional[UploadFile] = File(None),
     req: ImageEnhanceRequest = Depends(ImageEnhanceRequest.as_form),
     accept: str = Header(None),
     accept_query: str | None = Query(
         None, alias='accept',
         description="Parameter to override 'Accept' header, 'image/png' for output bytes")):
-    """\nImage Prompt\n
-    Image Prompt
-    A prompt-based image generation.
+    """\nImage Enhance\n
+    Enhance an image with optional input.
     Arguments:
-        enhance_input_image {UploadFile} -- Input image file
+        enhance_input_image {UploadFile} -- Input image file (optional)
         req {ImageEnhanceRequest} -- Request body
         accept {str} -- Accept header
         accept_query {str} -- Parameter to override 'Accept' header, 'image/png' for output bytes
@@ -177,9 +181,9 @@ def img_prompt(
 
 
 @secure_router.post(
-        path="/v1/tools/describe-image",
-        response_model=DescribeImageResponse,
-        tags=["GenerateV1"])
+    path="/v1/tools/describe-image",
+    response_model=DescribeImageResponse,
+    tags=["GenerateV1"])
 def describe_image(
     image: UploadFile,
     image_type: DescribeImageType = Query(
@@ -205,11 +209,11 @@ def describe_image(
 
 
 @secure_router.post(
-        path="/v1/generation/stop",
-        response_model=StopResponse,
-        description="Job stopping",
-        tags=["Default"])
+    path="/v1/generation/stop",
+    response_model=StopResponse,
+    description="Job stopping",
+    tags=["Default"])
 def stop():
     """Interrupt worker"""
     stop_worker()
-    return StopResponse(msg="success")
+    return StopResponse(msg="success")