comic/engine.py

"""generate_comic — the full idea -> finished comic pipeline.

It is a GENERATOR that yields progress events as it goes, so the Gradio app can show
live status and stream panels onto the page as they render (rather than freezing for
the whole ~minute of generation). Stages:

  1. WRITER bible call: safety gate + story bible. If refused -> a 'refused' event, done.
  2. WRITER panel calls: 5 batches of 2 pages each, each fed a recap of prior panels
     for continuity. Yields a 'panels' event per batch.
  3. ARTIST renders all 20 panels, one at a time, yielding an 'image' event each.
  4. 'done' event with the finished Comic.

The pipeline is backend-agnostic (mock or modal) via make_backends(). Errors in a
single panel render are caught so one bad image never sinks the whole comic.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Iterator, Optional

from .backends import make_backends, WriterBackend, ArtistBackend
from .schema import Comic, ComicBible, Panel, TOTAL_PANELS
from . import writer as W
from .imaging import build_image_prompt, panel_seed

# Panels rendered per GPU pass. 4 fits klein-9B at 832x576 on an H100 with headroom;
# the serve endpoint falls back to serial if a batch ever OOMs.
RENDER_BATCH = 4


@dataclass
class GenerateEvent:
    """One step of progress. `kind` drives the UI; payload fields are kind-specific."""
    kind: str                       # status|bible|refused|panels|image|done|error
    message: str = ""
    comic: Optional[Comic] = None   # carried from 'bible' onward (mutated in place)
    panel: Optional[Panel] = None   # for 'image' events
    progress: float = 0.0           # 0..1 coarse progress for a progress bar


def _retry_json(fn, parse, attempts: int = 2):
    """Call fn() -> reply, parse it; retry once on a parse/JSON failure."""
    last = None
    for _ in range(attempts):
        reply = fn()
        try:
            return parse(reply)
        except Exception as e:  # noqa: BLE001 - surface after retries
            last = e
    raise last if last else RuntimeError("generation failed")


def generate_comic(
    idea: str,
    writer: WriterBackend | None = None,
    artist: ArtistBackend | None = None,
    backend: str | None = None,
) -> Iterator[GenerateEvent]:
    """Yield GenerateEvents from raw idea to a fully rendered Comic."""
    if writer is None or artist is None:
        w, a = make_backends(backend)
        writer = writer or w
        artist = artist or a

    idea = (idea or "").strip()
    if not idea:
        yield GenerateEvent("error", "Please describe the comic you want.")
        return

    # ── 1. bible + safety gate ────────────────────────────────────────────────
    yield GenerateEvent("status", "Reading your idea and planning the story…", progress=0.02)
    try:
        bible: ComicBible = _retry_json(
            lambda: writer.chat(W.build_bible_messages(idea)),
            W.parse_bible,
        )
    except Exception as e:  # noqa: BLE001
        yield GenerateEvent("error", f"Couldn't plan the comic ({type(e).__name__}). Try again.")
        return

    if not bible.approved:
        reason = bible.refusal_reason or "That request can't be turned into a comic."
        yield GenerateEvent("refused", reason)
        return

    comic = Comic(bible=bible)
    yield GenerateEvent("bible", f"“{bible.title}” — {bible.logline}", comic=comic, progress=0.1)

    # ── 2. panel script, batched, with running recap for continuity ───────────
    written: list[Panel] = []
    page_batches = W.batches()
    n_batches = len(page_batches)
    for bi, pages in enumerate(page_batches):
        yield GenerateEvent(
            "status",
            f"Writing pages {pages[0]}–{pages[-1]} of {len(comic.bible.pages)}…",
            comic=comic,
            progress=0.1 + 0.3 * (bi / n_batches),
        )
        recap = W.recap_from_panels(written)
        try:
            panels = _retry_json(
                lambda: writer.chat(W.build_panel_messages(bible, pages, recap)),
                lambda r: W.parse_panels(r, pages),
            )
        except Exception as e:  # noqa: BLE001
            yield GenerateEvent("error", f"Story writing failed on pages {pages} ({type(e).__name__}).")
            return
        # Assemble each panel's image prompt now (deterministic, no model call).
        for p in panels:
            p.image_prompt = build_image_prompt(p, bible)
        written.extend(panels)
        comic.panels = sorted(written, key=lambda x: x.index)
        yield GenerateEvent("panels", f"Pages {pages[0]}–{pages[-1]} scripted.",
                            comic=comic, progress=0.1 + 0.3 * ((bi + 1) / n_batches))

    # ── 3. render every panel (batched through the GPU for throughput) ─────────
    ordered = sorted(comic.panels, key=lambda x: x.index)
    total = len(ordered) or TOTAL_PANELS
    done = 0
    for start in range(0, len(ordered), RENDER_BATCH):
        chunk = ordered[start:start + RENDER_BATCH]
        yield GenerateEvent(
            "status",
            f"Illustrating panels {start + 1}–{start + len(chunk)} of {total}…",
            comic=comic, progress=0.4 + 0.6 * (start / total),
        )
        prompts = [p.image_prompt for p in chunk]
        seeds = [panel_seed(bible, p) for p in chunk]
        try:
            images = artist.render_batch(prompts, seeds)
        except Exception as e:  # noqa: BLE001 - a batch failure must not sink the comic
            images = [None] * len(chunk)
            yield GenerateEvent("status", f"A render batch hiccupped ({type(e).__name__}); continuing…",
                                comic=comic)
        for panel, img in zip(chunk, images):
            panel.image = img
            done += 1
            yield GenerateEvent("image", f"Panel {done} ready.", comic=comic, panel=panel,
                                progress=0.4 + 0.6 * (done / total))

    yield GenerateEvent("done", f"“{bible.title}” is ready.", comic=comic, progress=1.0)