Add aspect ratio and image size options to image generation; implement chat-based image generation handling

Co-authored-by: Copilot <copilot@github.com>
This commit is contained in:
2026-04-27 19:28:04 +02:00
parent 17ae8d9477
commit 58c2cb4490
6 changed files with 211 additions and 18 deletions
+56 -18
View File
@@ -58,31 +58,69 @@ async def generate_image(
_: dict = Depends(get_current_user),
) -> ImageResponse:
"""Generate images from a text prompt."""
# Detect if model uses chat completions (FLUX, GPT-5 Image Mini) vs /images/generations (DALL-E)
chat_models = {"black-forest-labs/flux.2-klein-4b",
"openai/gpt-5-image-mini"}
is_chat_model = body.model.lower() in {m.lower() for m in chat_models} or \
any(m in body.model.lower() for m in ["flux", "gpt-5-image-mini"])
try:
result = await openrouter.generate_image(
model=body.model,
prompt=body.prompt,
n=body.n,
size=body.size,
)
if is_chat_model:
image_config = {}
if body.aspect_ratio:
image_config["aspect_ratio"] = body.aspect_ratio
if body.image_size:
image_config["image_size"] = body.image_size
result = await openrouter.generate_image_chat(
model=body.model,
prompt=body.prompt,
modalities=[
"image", "text"] if "gpt-5-image-mini" in body.model.lower() else ["image"],
image_config=image_config if image_config else None,
)
else:
result = await openrouter.generate_image(
model=body.model,
prompt=body.prompt,
n=body.n,
size=body.size,
)
except Exception as exc:
raise HTTPException(
status_code=status.HTTP_502_BAD_GATEWAY, detail=f"OpenRouter error: {exc}")
try:
images = [
ImageResult(
url=item.get("url"),
b64_json=item.get("b64_json"),
revised_prompt=item.get("revised_prompt"),
if is_chat_model:
# Chat completions response: choices[0].message.images[].image_url.url
images = []
message = result.get("choices", [{}])[0].get("message", {})
for item in message.get("images", []):
img_url = item.get("image_url", {}).get("url")
images.append(ImageResult(
url=img_url,
b64_json=None,
revised_prompt=message.get("content"),
))
return ImageResponse(
id=result.get("id", ""),
model=result.get("model", body.model),
images=images,
)
else:
# /images/generations response: data[].url
images = [
ImageResult(
url=item.get("url"),
b64_json=item.get("b64_json"),
revised_prompt=item.get("revised_prompt"),
)
for item in result.get("data", [])
]
return ImageResponse(
id=result.get("id", ""),
model=result.get("model", body.model),
images=images,
)
for item in result.get("data", [])
]
return ImageResponse(
id=result.get("id", ""),
model=result.get("model", body.model),
images=images,
)
except (KeyError, TypeError) as exc:
raise HTTPException(status_code=status.HTTP_502_BAD_GATEWAY,
detail=f"Unexpected response format: {exc}")