image_generator.py

"""
image_generator.py — Generates theatrical scene images via FLUX.2-pro on OpenRouter.

OpenRouter image generation uses the /api/v1/chat/completions endpoint
with modalities: ["image"], returning base64-encoded images.
"""

import os
import base64
import requests
from pathlib import Path
from PIL import Image
import io

_PROMPT_PREFIX = "theatrical stage scene, dramatic lighting, cinematic, "
_MODEL = "black-forest-labs/flux.2-pro"
_OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions"


class ImageGenerator:
    """Generates 16:9 theatrical scene images via FLUX on OpenRouter."""

    def __init__(self, api_key: str):
        self.api_key = api_key
        if not self.api_key:
            print("[ImageGenerator] WARNING: api_key is empty — image generation will fail.")

    def generate(self, prompt: str, output_path: str) -> str:
        """
        Generate a 16:9 theatrical scene image and save as PNG to output_path.

        Returns output_path on success.
        Raises RuntimeError on failure.
        """
        if not self.api_key:
            raise RuntimeError("OpenRouter API key not configured.")

        full_prompt = _PROMPT_PREFIX + prompt

        Path(output_path).parent.mkdir(parents=True, exist_ok=True)

        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }
        payload = {
            "model": _MODEL,
            "messages": [
                {"role": "user", "content": full_prompt}
            ],
            "modalities": ["image"],
            "image_config": {
                "aspect_ratio": "16:9",
            },
        }

        response = requests.post(_OPENROUTER_URL, headers=headers, json=payload, timeout=120)

        if response.status_code != 200:
            raise RuntimeError(
                f"OpenRouter image API error {response.status_code}: {response.text[:300]}"
            )

        data = response.json()

        # Extract image from response
        # OpenRouter returns images in choices[0].message.images or as base64 data URL
        image_bytes = None
        try:
            message = data["choices"][0]["message"]

            # Try message.images list first
            if message.get("images"):
                img_url = message["images"][0]["image_url"]["url"]
            # Fall back to message.content as data URL
            elif isinstance(message.get("content"), str) and message["content"].startswith("data:image"):
                img_url = message["content"]
            else:
                raise RuntimeError(f"No image found in response: {str(data)[:400]}")

            # Decode base64 data URL
            if img_url.startswith("data:image"):
                header, b64data = img_url.split(",", 1)
                image_bytes = base64.b64decode(b64data)
            else:
                # It's a URL — fetch it
                r = requests.get(img_url, timeout=60)
                r.raise_for_status()
                image_bytes = r.content

        except (KeyError, IndexError) as e:
            raise RuntimeError(f"Unexpected response format: {e}\n{str(data)[:400]}")

        # Save as PNG
        img = Image.open(io.BytesIO(image_bytes))
        img.save(output_path, format="PNG")
        return output_path