Skip to content

Commit

Permalink
runner: Multiple outputs for text/image-to-image
Browse files Browse the repository at this point in the history
  • Loading branch information
yondonfu committed Feb 7, 2024
1 parent fd12195 commit 87e4f48
Show file tree
Hide file tree
Showing 5 changed files with 52 additions and 32 deletions.
2 changes: 2 additions & 0 deletions runner/app/routes/image_to_image.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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}")
Expand Down
1 change: 1 addition & 0 deletions runner/app/routes/text_to_image.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}}
Expand Down
10 changes: 10 additions & 0 deletions runner/openapi.json
Original file line number Diff line number Diff line change
Expand Up @@ -248,6 +248,11 @@
"seed": {
"type": "integer",
"title": "Seed"
},
"num_images_per_prompt": {
"type": "integer",
"title": "Num Images Per Prompt",
"default": 1
}
},
"type": "object",
Expand Down Expand Up @@ -402,6 +407,11 @@
"seed": {
"type": "integer",
"title": "Seed"
},
"num_images_per_prompt": {
"type": "integer",
"title": "Num Images Per Prompt",
"default": 1
}
},
"type": "object",
Expand Down
5 changes: 5 additions & 0 deletions worker/multipart.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
66 changes: 34 additions & 32 deletions worker/runner.gen.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit 87e4f48

Please sign in to comment.