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

Co-authored-by: Copilot <copilot@github.com>

backend/app/models/ai.py

@@ -53,6 +53,8 @@ class ImageRequest(BaseModel):
     prompt: str
     n: int = 1
     size: str = "1024x1024"
+    aspect_ratio: str | None = None  # e.g. "1:1", "16:9", "9:16"
+    image_size: str | None = None  # e.g. "0.5K", "1K", "2K", "4K"
 
 
 class ImageResult(BaseModel):

backend/app/routers/generate.py

@@ -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}")

backend/app/services/openrouter.py

@@ -139,3 +139,34 @@ async def poll_video_status(polling_url: str) -> dict[str, Any]:
         response = await client.send(resp)
         response.raise_for_status()
         return response.json()
+
+
+async def generate_image_chat(
+    model: str,
+    prompt: str,
+    modalities: list[str] | None = None,
+    image_config: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    """Request image generation via Chat Completions with modalities.
+
+    Used by models like FLUX.2 Klein 4B and GPT-5 Image Mini that output
+    images through the chat completions endpoint rather than /images/generations.
+    """
+    base_url = os.getenv("OPENROUTER_BASE_URL", OPENROUTER_BASE_URL)
+    if modalities is None:
+        # Image-only models (FLUX) vs multimodal (GPT-5 Image Mini)
+        modalities = ["image"]
+    payload: dict[str, Any] = {
+        "model": model,
+        "messages": [{"role": "user", "content": prompt}],
+        "modalities": modalities,
+    }
+    if image_config:
+        payload["image_config"] = image_config
+    async with httpx.AsyncClient(timeout=120) as client:
+        resp = client.build_request(
+            "POST", f"{base_url}/chat/completions", headers=_headers(), json=payload
+        )
+        response = await client.send(resp)
+        response.raise_for_status()
+        return response.json()