tests/probe_comic.py

"""Manual probe: render a few comic prompts on the live FLUX endpoint, print
latency, and save the PNGs for visual sign-off on the comic style.

Usage:  set FLUX_URL and FLUX_TOKEN, then  python tests/probe_comic.py
Saved samples land in  tests/comic_samples/ .
"""
import base64
import os
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from game import comic  # noqa: E402
from game.prompts import COMIC_STYLE  # noqa: E402

if not comic.flux_available():
    print("FLUX_URL not set — nothing to probe (comic falls back to no overlay).")
    sys.exit(1)

OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "comic_samples")
os.makedirs(OUT, exist_ok=True)

PROMPTS = [
    COMIC_STYLE + ". Panel 1: a chibi office worker sweating, wide-eyed, hands "
    "on his face at a laptop. Panel 2: the whole office turning to stare at him.",
    COMIC_STYLE + ". Single panel: a chibi worker hiding under her desk hugging "
    "a coffee mug, forced nervous smile.",
    COMIC_STYLE + ". Panel 1: a fat envelope of cash on a desk. Panel 2: a boss "
    "raising one eyebrow, arms crossed. Panel 3: the envelope being slid back.",
]

failures = 0
for i, prompt in enumerate(PROMPTS, 1):
    t0 = time.time()
    img = comic.generate_comic(prompt)
    ms = int((time.time() - t0) * 1000)
    if img:
        path = os.path.join(OUT, f"comic_{i}.png")
        with open(path, "wb") as f:
            f.write(base64.b64decode(img))
        print(f"OK   panel {i}  {ms:6d}ms  {len(img):>8}b  -> {path}")
    else:
        failures += 1
        print(f"FAIL panel {i}  {ms:6d}ms  (no image — would fall back to no overlay)")

sys.exit(1 if failures else 0)