Skip to content

Commit

Permalink
process text-to-image requested image count sequentially
Browse files Browse the repository at this point in the history
  • Loading branch information
ad-astra-video committed May 20, 2024
1 parent c05cccf commit 5ab6f10
Showing 1 changed file with 25 additions and 17 deletions.
42 changes: 25 additions & 17 deletions runner/app/routes/text_to_image.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
from app.routes.util import image_to_data_url, ImageResponse, HTTPError, http_error
import logging
import random
import os
import os, json

router = APIRouter()

Expand Down Expand Up @@ -59,24 +59,32 @@ async def text_to_image(

if params.seed is None:
params.seed = random.randint(0, 2**32 - 1)
if params.num_images_per_prompt > 1:
params.seed = [
i for i in range(params.seed, params.seed + params.num_images_per_prompt)
]

try:
images, has_nsfw_concept = pipeline(**params.model_dump())
except Exception as e:
logger.error(f"TextToImagePipeline error: {e}")
logger.exception(e)
return JSONResponse(
status_code=500, content=http_error("TextToImagePipeline error")
)
if params.num_images_per_prompt > 1:
params.seed = [
i for i in range(params.seed, params.seed + params.num_images_per_prompt)
]

if not isinstance(params.seed, list):
params.seed = [params.seed]

num_images_per_prompt = params.num_images_per_prompt
params.num_images_per_prompt = 1
images = []
has_nsfw_concept = []
seeds = params.seed
if not isinstance(seeds, list):
seeds = [seeds]

for i in range(num_images_per_prompt):
try:
params.seed = [seeds[i]] #pass one seed at a time
imgs, nsfw_check = pipeline(**params.model_dump())
images.extend(imgs)
has_nsfw_concept.extend(nsfw_check)
except Exception as e:
logger.error(f"TextToImagePipeline error: {e}")
logger.exception(e)
return JSONResponse(
status_code=500, content=http_error("TextToImagePipeline error")
)

output_images = []
for img, sd, is_nsfw in zip(images, seeds, has_nsfw_concept):
# TODO: Return None once Go codegen tool supports optional properties
Expand Down

0 comments on commit 5ab6f10

Please sign in to comment.