runner: Multiple outputs for text/image-to-image

Author: yondonfu

runner/app/routes/image_to_image.py

@@ -31,6 +31,7 @@ async def image_to_image(
     guidance_scale: Annotated[float, Form()] = 7.5,
     negative_prompt: Annotated[str, Form()] = "",
     seed: Annotated[int, Form()] = None,
+    num_images_per_prompt: Annotated[int, Form()] = 1,
     pipeline: ImageToImagePipeline = Depends(get_pipeline),
 ):
     if model_id != "" and model_id != pipeline.model_id:
@@ -49,6 +50,7 @@ async def image_to_image(
             guidance_scale=guidance_scale,
             negative_prompt=negative_prompt,
             seed=seed,
+            num_images_per_prompt=num_images_per_prompt,
         )
     except Exception as e:
         logger.error(f"ImageToImagePipeline error: {e}")

runner/app/routes/text_to_image.py

@@ -21,6 +21,7 @@ class TextToImageParams(BaseModel):
     guidance_scale: float = 7.5
     negative_prompt: str = ""
     seed: int = None
+    num_images_per_prompt: int = 1
 
 
 responses = {400: {"model": HTTPError}, 500: {"model": HTTPError}}

runner/openapi.json

@@ -248,6 +248,11 @@
           "seed": {
             "type": "integer",
             "title": "Seed"
+          },
+          "num_images_per_prompt": {
+            "type": "integer",
+            "title": "Num Images Per Prompt",
+            "default": 1
           }
         },
         "type": "object",
@@ -402,6 +407,11 @@
           "seed": {
             "type": "integer",
             "title": "Seed"
+          },
+          "num_images_per_prompt": {
+            "type": "integer",
+            "title": "Num Images Per Prompt",
+            "default": 1
           }
         },
         "type": "object",

worker/multipart.go

@@ -54,6 +54,11 @@ func NewImageToImageMultipartWriter(w io.Writer, req ImageToImageMultipartReques
 			return nil, err
 		}
 	}
+	if req.NumImagesPerPrompt != nil {
+		if err := mw.WriteField("num_images_per_prompt", strconv.Itoa(*req.NumImagesPerPrompt)); err != nil {
+			return nil, err
+		}
+	}
 
 	if err := mw.Close(); err != nil {
 		return nil, err