image_engine.py

"""
Image Engine — Generates scene images via OpenRouter using black-forest-labs/flux.2-pro.
Saves images as compressed JPEG files to avoid huge base64 data URIs.
"""

import os
import re
import io
import base64
import requests
from PIL import Image
from dotenv import load_dotenv

load_dotenv()

OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY", "")
MODEL_ID = "black-forest-labs/flux.2-pro"
OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions"

# Directory to store generated scene images
IMG_DIR = os.path.join(os.path.dirname(__file__), "generated_scenes")
os.makedirs(IMG_DIR, exist_ok=True)

_img_counter = 0


def _save_as_jpeg(data_uri: str) -> str:
    """Convert a base64 data URI (PNG) to a compressed JPEG file. Returns file path."""
    global _img_counter
    _img_counter += 1

    # Strip data URI prefix
    b64_data = data_uri.split(",", 1)[1] if "," in data_uri else data_uri
    raw = base64.b64decode(b64_data)

    img = Image.open(io.BytesIO(raw)).convert("RGB")
    path = os.path.join(IMG_DIR, f"scene_{_img_counter}.jpg")
    img.save(path, "JPEG", quality=72)

    size_kb = os.path.getsize(path) / 1024
    print(f"[ImageEngine] Saved JPEG ({size_kb:.0f} KB): {path}")
    return path


def generate_scene_image(prompt: str) -> str | None:
    """
    Generate a scene image from a text prompt via OpenRouter chat/completions.
    Returns a local file path to a JPEG or None on failure.
    """
    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type": "application/json",
    }

    full_prompt = (
        f"A cinematic wide-angle theatrical stage scene, dramatic lighting, "
        f"rich colors, film still quality: {prompt}"
    )

    payload = {
        "model": MODEL_ID,
        "messages": [
            {"role": "user", "content": full_prompt}
        ],
        "modalities": ["image"],
    }

    try:
        print(f"[ImageEngine] Generating image: {prompt[:60]}...")
        resp = requests.post(
            OPENROUTER_URL,
            headers=headers,
            json=payload,
            timeout=60,
        )

        if resp.status_code != 200:
            print(f"[ImageEngine] API Error ({resp.status_code}): {resp.text[:200]}")
            return None

        data = resp.json()
        message = data.get("choices", [{}])[0].get("message", {})

        # Flux returns images in message.images[] array
        images = message.get("images", [])
        if images:
            img = images[0]
            if isinstance(img, dict):
                url = img.get("image_url", {}).get("url") or img.get("url")
                if url and url.startswith("data:image"):
                    return _save_as_jpeg(url)
                elif url:
                    return url  # External URL, use directly

        # Fallback: check content field
        content = message.get("content")
        if content:
            if content.startswith("data:image"):
                return _save_as_jpeg(content)
            md_match = re.search(r'!\[.*?\]\((.*?)\)', content)
            if md_match:
                return md_match.group(1)

        print(f"[ImageEngine] Could not extract image from response.")
        return None

    except Exception as e:
        print(f"[ImageEngine] Error: {e}")
        return None