game/comic.py

"""Modal FLUX client for comic-panel image generation.

If FLUX_URL is set, POSTs the image prompt to the Modal FLUX endpoint and
returns a base64 PNG. Otherwise — or on any failure/timeout — returns None,
and the frontend renders the composed chibi-comic fallback instead. Mirrors the
warm / cold-timeout pattern in llm.py.
"""
from __future__ import annotations

import os
import threading
import time

import requests

from .trace import trace

# 30s warm default: on a hosted Space a slightly-cold / loaded FLUX container
# needs more than 12s, else panels silently no-op. Warm renders are ~3-6s.
FLUX_TIMEOUT = float(os.environ.get("BDS_FLUX_TIMEOUT", "30"))
FLUX_COLD_TIMEOUT = float(os.environ.get("BDS_FLUX_COLD_TIMEOUT", "90"))
_warmed = False


def flux_available() -> bool:
    return bool(os.environ.get("FLUX_URL"))


def _headers() -> dict:
    return {"Authorization": "Bearer " + os.environ.get("FLUX_TOKEN", "")}


def warm() -> None:
    """Boot the FLUX container in the background (page load / new game)."""
    if not flux_available() or _warmed:
        return
    threading.Thread(target=_warm_ping, daemon=True).start()


def _warm_ping() -> None:
    try:
        requests.post(os.environ["FLUX_URL"], json={"warmup": True},
                      headers=_headers(), timeout=FLUX_COLD_TIMEOUT + 60)
        trace("flux", "warmup ping done")
    except requests.RequestException:
        pass


def generate_comic(prompt: str) -> str | None:
    """Return a base64 PNG for the comic, or None → composed fallback."""
    if not flux_available() or not prompt:
        return None
    global _warmed
    timeout = FLUX_TIMEOUT if _warmed else FLUX_COLD_TIMEOUT
    t0 = time.time()
    try:
        resp = requests.post(os.environ["FLUX_URL"], json={"prompt": prompt},
                             headers=_headers(), timeout=timeout)
        ms = int((time.time() - t0) * 1000)
        if resp.status_code == 200:
            body = resp.json()
            img = body.get("image_b64")
            if img:
                _warmed = True
                trace("flux", f"comic LIVE {ms}ms ({len(img)}b)")
                return img
            trace("flux", f"comic not-ok {ms}ms: "
                          f"{str(body.get('error', body))[:120]}")
        else:
            trace("flux", f"comic HTTP {resp.status_code} {ms}ms")
    except requests.Timeout:
        trace("flux", f"comic TIMEOUT after {timeout}s")
    except requests.RequestException as exc:
        trace("flux", f"comic transport error: {str(exc)[:120]}")
    return None