diffu_page/render.py

"""Render stage — turn one line's text into an ink-on-transparent raster for compositing.

:class:`LineRenderer` is the seam between "what to write" and "how it looks". Two implementations:

- :class:`PlaceholderRenderer` draws gray text with OpenCV (no model, no GPU, no torch) so the
  layout + compositor are runnable and testable *today*.
- :class:`DiffuRenderer` wraps :func:`diffu.generate.generate_line` (lazy torch import) — one writer
  style per page, enforced by taking the style reference at construction.

Output convention (both): a ``BGRA`` ``uint8`` array whose **alpha channel is the ink coverage**
(255 = full ink, 0 = paper). BGR is left neutral — the ink *colour* is chosen later by the
compositor from the local paper, so blending stays adaptive.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any, Protocol

import cv2
import numpy as np
from pydantic import BaseModel

if TYPE_CHECKING:
    import torch
    from PIL.Image import Image

    from diffu.config import Config
    from diffu.model import Diffu

    from .layout import Layout, LineBox


class LineJob(BaseModel):
    """One line to render: its text, the target ink height (px), and the writer-style reference path."""

    model_config = {"frozen": True}

    text: str
    height: int
    style_path: str


class LineRenderer(Protocol):
    """Renders one text line to a BGRA ink raster at approximately ``height`` pixels tall."""

    def render(self, text: str, height: int) -> np.ndarray: ...


class PlaceholderRenderer:
    """Gray Hershey-font text → BGRA ink. For developing/testing the page pipeline without a model."""

    def __init__(self, font: int = cv2.FONT_HERSHEY_SIMPLEX) -> None:
        self.font = font

    def render(self, text: str, height: int) -> np.ndarray:
        text = text or " "
        thickness = max(1, height // 22)
        scale = cv2.getFontScaleFromHeight(self.font, max(8, int(height * 0.62)), thickness)
        (tw, th), _ = cv2.getTextSize(text, self.font, scale, thickness)
        pad = max(4, int(height * 0.25))
        canvas = np.zeros((height + 2 * pad, tw + 2 * pad), np.uint8)  # alpha/coverage
        baseline_y = pad + th
        cv2.putText(canvas, text, (pad, baseline_y), self.font, scale, 255, thickness, cv2.LINE_AA)
        bgra = np.zeros((canvas.shape[0], canvas.shape[1], 4), np.uint8)
        bgra[..., 3] = canvas
        return bgra

    def measure_width(self, text: str, height: int) -> int:
        """Predicted rendered width of ``text`` at ``height`` — matches :meth:`render` exactly."""
        text = text or " "
        thickness = max(1, height // 22)
        scale = cv2.getFontScaleFromHeight(self.font, max(8, int(height * 0.62)), thickness)
        (tw, _), _ = cv2.getTextSize(text, self.font, scale, thickness)
        return tw + 2 * max(4, int(height * 0.25))


def _assert_conditioner_loaded(model: Diffu, ckpt: str) -> None:
    """Fail fast if a checkpoint's TEXT-CONDITIONER (``glyph_content``) doesn't fit this code's architecture.

    ``load_checkpoint`` is non-strict, so a checkpoint trained with a different conditioner silently leaves
    that module at its RANDOM init — the model then generates confident-looking GIBBERISH that no metric on
    the output image catches (this masked a whole debugging session). Raise only on a genuine architecture
    change: much of the conditioner absent AND the checkpoint carrying differently-named conditioner
    weights — never the benign frozen-and-legitimately-absent case.
    """
    if not ckpt.endswith(".safetensors"):
        return
    from safetensors import safe_open

    with safe_open(ckpt, framework="pt") as handle:
        ckpt_keys = set(handle.keys())
    own = set(model.state_dict())
    cond = [k for k in own if k.startswith("glyph_content.")]
    missing = [k for k in cond if k not in ckpt_keys]
    renamed = [k for k in ckpt_keys - own if k.startswith("glyph_content.")]
    if cond and len(missing) > 0.5 * len(cond) and renamed:
        raise RuntimeError(
            f"checkpoint architecture mismatch: {len(missing)}/{len(cond)} text-conditioner (glyph_content) "
            f"weights are absent and {len(renamed)} differently-named ones are present, so the conditioner "
            f"stays random and the model generates gibberish. This checkpoint was trained with a different "
            f"model version — generate from a code-matching checkpoint instead.\n  checkpoint: {ckpt}"
        )


class DiffuRenderer:
    """Wrap Diffu's line generator: (text) → BGRA ink, one fixed writer style for the whole page.

    Heavy deps (torch, the model, the checkpoint) are imported and built lazily on first ``render``,
    so importing :mod:`diffu_page` and running the layout/compositor needs only OpenCV + numpy.
    """

    def __init__(
        self,
        ckpt: str,
        style_path: str,
        *,
        cfg: Config | None = None,
        num_steps: int = 24,
        cfg_scale: float = 3.0,
        device: str | None = None,
        compile_model: bool = False,
    ) -> None:
        self.ckpt = ckpt
        self.style_path = style_path
        self.cfg = cfg
        self.num_steps = num_steps
        self.cfg_scale = cfg_scale
        self.device = device
        self.compile_model = compile_model
        self._model: Diffu | None = None
        self._style: torch.Tensor | None = None

    def _ensure(self) -> None:
        if self._model is not None:
            return
        import torch  # lazy

        from diffu.config import Config
        from diffu.generate import load_checkpoint, load_style
        from diffu.model import Diffu

        self.cfg = self.cfg or Config()
        self.device = self.device or ("cuda" if torch.cuda.is_available() else "cpu")
        model = Diffu(self.cfg).to(self.device).eval()
        if self.ckpt:
            load_checkpoint(model, self.ckpt)
            _assert_conditioner_loaded(model, self.ckpt)  # fail fast: a mismatched arch generates gibberish
        if self.compile_model:
            model.backbone.compile_blocks()  # a page = many lines → compile cold-start amortizes well
        self._model = model
        self._style = load_style(self.style_path).unsqueeze(0).to(self.device)  # [1,3,S,S]

    def render(self, text: str, height: int) -> np.ndarray:
        import torch  # lazy

        from diffu.generate import generate_line

        self._ensure()
        if self._model is None or self._style is None or self.cfg is None:
            raise RuntimeError("DiffuRenderer not initialised")  # _ensure guarantees these
        with torch.no_grad():
            pil = generate_line(
                self._model,
                text or " ",
                self._style,
                cfg=self.cfg,
                num_steps=self.num_steps,
                cfg_scale=self.cfg_scale,
            )
        return _ink_from_pil(pil, height)

    def measure_width(self, text: str, height: int) -> int:
        """Predicted width at ``height`` from the model's content-aware natural width (no denoise)."""
        from diffu.generate import natural_width

        self._ensure()
        if self._model is None or self.cfg is None:
            raise RuntimeError("DiffuRenderer not initialised")
        w = natural_width(self._model.guidance_renderer, text or " ", self.cfg)
        return int(w * height / self.cfg.data.line_height)


class ErukuRenderer:
    """A :class:`LineRenderer` backed by Eruku — a public autoregressive styled-text generator.

    Lets the page pipeline run with REAL handwriting without a Diffu checkpoint (any line generator
    that satisfies :class:`LineRenderer` drops in here). Reuses the project's torch + transformers
    (lazy import); one style reference = one writer per page.

    Caveat: Eruku is English/modern (IAM/CVL/RIMES/FontSquare), NOT Swedish-historical — use it to
    validate the *pipeline*, not to produce domain-correct data; å ä ö may render poorly. A GPU is
    needed in practice (autoregressive generation per line).
    """

    def __init__(
        self,
        style_path: str,
        *,
        cfg_scale: float = 1.25,
        model_id: str = "blowing-up-groundhogs/eruku",
        device: str | None = None,
    ) -> None:
        self.style_path = style_path
        self.cfg_scale = cfg_scale
        self.model_id = model_id
        self.device = device
        self._model: Any = None
        self._style: Image | None = None

    def _ensure(self) -> None:
        if self._model is not None:
            return
        import torch  # lazy
        from PIL import Image
        from transformers import AutoModel

        self.device = self.device or ("cuda" if torch.cuda.is_available() else "cpu")
        self._model = AutoModel.from_pretrained(self.model_id, trust_remote_code=True).to(self.device).eval()
        self._style = Image.open(self.style_path).convert("RGB")

    def render(self, text: str, height: int) -> np.ndarray:
        import torch  # lazy

        self._ensure()
        if self._style is None:
            raise RuntimeError("ErukuRenderer not initialised")  # _ensure guarantees this
        with torch.inference_mode():
            pil = self._model.generate_handwriting(
                style_image=self._style, gen_text=text or " ", style_text="", cfg_scale=self.cfg_scale
            )
        return _ink_from_pil(pil, height)


def _ink_from_pil(pil: Image, height: int) -> np.ndarray:
    """A generated line image -> BGRA, resized to ``height`` (aspect kept).

    The compositor never re-renders the ink — it alpha-OVERs the model's *real pixels*. So we carry:

    - **BGR** = the model's actual line patch, verbatim (ink + its paper), resized.
    - **alpha** (A) = a soft ink matte (see :func:`_ink_matte`) that keeps the strokes and maps the line's
      own grey writing-region wash — and any white padding around a real crop — to 0, so only real ink
      survives and sits directly on the page paper (no off-paper box or ``[ ]`` frame around the line).
    """
    bgr = cv2.cvtColor(np.asarray(pil.convert("RGB")), cv2.COLOR_RGB2BGR).astype(np.float32)
    gray = cv2.cvtColor(bgr.astype(np.uint8), cv2.COLOR_BGR2GRAY).astype(np.float32)
    alpha = _ink_matte(gray)
    h, w = gray.shape
    if h != height:
        nw = max(1, int(w * height / h))
        bgr = cv2.resize(bgr, (nw, height), interpolation=cv2.INTER_AREA)
        alpha = cv2.resize(alpha, (nw, height), interpolation=cv2.INTER_AREA)
    bgra = np.empty((*alpha.shape, 4), np.uint8)
    bgra[..., :3] = np.clip(bgr, 0, 255).astype(np.uint8)
    bgra[..., 3] = (alpha * 255.0).astype(np.uint8)
    return bgra


def _ink_matte(gray: np.ndarray) -> np.ndarray:
    """Soft ink coverage in ``[0, 1]`` via LOCAL background subtraction — keeps the full strokes the model
    drew while flattening its grey writing-region wash to 0.

    A flat brightness threshold can't separate the lighter parts of a stroke from the grey wash (both are
    mid-grey), so it either keeps the wash as a box OR eats the strokes down to broken fragments. Instead,
    estimate the LOCAL background and take ``bg - gray``: a pixel darker than its surroundings (the whole
    stroke, core + lighter edges) scores high, while the flat uniform wash (== its own local background)
    scores ~0. A small floor removes the faint wash residual.

    Two details keep this clean on *real* line crops (grey writing-region on white padding):

    - **Median, not morphological close, for the background.** Grayscale close dilates first, so at a step
      edge (the grey wash meeting the white padding at a line's start/end) it drags the bright padding
      *into* the wash, and ``bg - gray`` lights up the whole region perimeter -> an off-paper ``[ ]`` frame
      around every seated line. A median is robust to that step, so no frame. It also collapses solid dark
      BANDS (a scan smudge reads as its own background -> ~0 ink) that close would keep.
    - **Flatten the white padding to the wash first.** With the padding still pure white, even the median
      window straddling the wash/padding boundary leaves a faint tick; mapping it to the local wash removes
      the step entirely. On model output (no white padding) this is a no-op."""
    g = gray.copy()
    body = gray[gray < 250.0]  # everything but the crop's bright white padding
    if body.size:
        g[gray >= 250.0] = float(
            np.percentile(body, 60)
        )  # flatten padding to the local wash: kill the step edge
    k = max(11, (g.shape[0] // 3) | 1)  # window a bit wider than a stroke, to see across it to the background
    bg = cv2.medianBlur(g.astype(np.uint8), k).astype(
        np.float32
    )  # median bg: robust to step edges + solid bands
    ink = bg - g  # how much darker than the LOCAL background -> the full stroke
    positive = ink[ink > 2.0]
    scale = float(np.percentile(positive, 90)) if positive.size else 30.0
    alpha = np.clip(ink / max(scale, 10.0), 0.0, 1.0)
    alpha = np.clip((alpha - 0.22) / 0.78, 0.0, 1.0)  # floor: drop the faint wash residual, keep the strokes
    return _suppress_blobs(alpha)


def _suppress_blobs(alpha: np.ndarray) -> np.ndarray:
    """Zero out solid high-fill components — source-scan smudges and the model's grey wash BANDS — while
    leaving handwriting untouched. A band is a filled block (large area, high bbox fill, tall); a stroke is
    a thin curve (low fill) or small, so the thresholds never catch real writing (verified on real crops)."""
    h = alpha.shape[0]
    mask = (alpha > 0.30).astype(np.uint8)
    count, labels, stats, _ = cv2.connectedComponentsWithStats(mask, connectivity=8)
    out = alpha.copy()
    for c in range(1, count):
        area = int(stats[c, cv2.CC_STAT_AREA])
        cw, ch = int(stats[c, cv2.CC_STAT_WIDTH]), int(stats[c, cv2.CC_STAT_HEIGHT])
        if area > 0.03 * alpha.size and area / max(1, cw * ch) > 0.5 and ch > 0.35 * h:
            out[labels == c] = 0.0
    return out


class BatchedDiffuRenderer:
    """Render every line of a whole job in one batched GPU pass — the page pipeline's main speedup.

    Per-line generation feeds the GPU one ~64px line at a time (it sits ~half-idle). Here the whole run's
    lines are width-bucketed and each bucket pushed through the model's batched ``generate`` in chunks of
    ``max_batch`` (optionally with fused/compiled blocks, which amortise over the few bucket widths), then
    ink-cropped. The model is built once; each line keeps its own writer style.
    """

    def __init__(
        self,
        ckpt: str,
        *,
        cfg: Config | None = None,
        num_steps: int = 24,
        cfg_scale: float = 3.0,
        device: str | None = None,
        max_batch: int = 48,
        bucket_px: int = 16,  # near-natural width: a coarse bucket makes the model spread/soften strokes
        compile_blocks: bool = True,
    ) -> None:
        self.ckpt = ckpt
        self.cfg = cfg
        self.num_steps = num_steps
        self.cfg_scale = cfg_scale
        self.device = device
        self.max_batch = max_batch
        self.bucket_px = bucket_px  # width quantisation: coarser -> fewer compiles, more padding waste
        self.compile_blocks = compile_blocks
        self._model: Diffu | None = None
        self._styles: dict[str, torch.Tensor] = {}

    def _ensure(self) -> None:
        if self._model is not None:
            return
        import torch  # lazy

        from diffu.config import Config
        from diffu.generate import load_checkpoint
        from diffu.model import Diffu

        self.cfg = self.cfg or Config()
        self.device = self.device or ("cuda" if torch.cuda.is_available() else "cpu")
        torch.set_float32_matmul_precision("high")  # TF32 matmuls — free on Ampere+, no quality cost
        model = Diffu(self.cfg).to(self.device).eval()
        if self.ckpt:
            load_checkpoint(model, self.ckpt)
        if self.compile_blocks and hasattr(model.backbone, "compile_blocks"):
            model.backbone.compile_blocks()  # fuse per-block glue; amortised over the few bucket widths
        self._model = model

    def _style(self, path: str) -> torch.Tensor:
        if path not in self._styles:
            from diffu.generate import load_style

            self._styles[path] = load_style(path).unsqueeze(0).to(self.device)  # [1, 3, S, S]
        return self._styles[path]

    def measure_width(self, text: str, height: int) -> int:
        """Predicted rendered width of ``text`` at ``height`` from the model's content-aware natural width
        (no denoise) — lets the layout wrap by real width before the batched render."""
        from diffu.generate import natural_width

        self._ensure()
        if self._model is None or self.cfg is None:
            raise RuntimeError("renderer not initialised")  # _ensure guarantees these
        w = natural_width(self._model.guidance_renderer, text or " ", self.cfg)
        return int(w * height / self.cfg.data.line_height)

    def render_jobs(self, jobs: list[LineJob]) -> list[np.ndarray]:
        """Render each job to a BGRA ink raster, width-bucketed and batched. Output aligns with ``jobs``."""
        import torch  # lazy

        from diffu.generate import ink_crop, natural_width, to_pil

        self._ensure()
        if self._model is None or self.cfg is None:
            raise RuntimeError("renderer not initialised")  # _ensure guarantees these
        ds = self.cfg.vae.downscale_factor
        latent_h = self.cfg.data.line_height // ds

        widths = [natural_width(self._model.guidance_renderer, j.text or " ", self.cfg) for j in jobs]
        bucket = [((w + self.bucket_px - 1) // self.bucket_px) * self.bucket_px for w in widths]
        order = sorted(range(len(jobs)), key=lambda i: bucket[i])  # group equal bucket widths together

        inks: list[np.ndarray | None] = [None] * len(jobs)
        i = 0
        while i < len(order):
            width = bucket[order[i]]
            group: list[int] = []
            while i < len(order) and bucket[order[i]] == width:
                group.append(order[i])
                i += 1
            latent_w = width // ds
            for start in range(0, len(group), self.max_batch):
                chunk = group[start : start + self.max_batch]
                texts = [jobs[k].text or " " for k in chunk]
                style = torch.cat([self._style(jobs[k].style_path) for k in chunk], dim=0)
                with torch.no_grad():
                    imgs = self._model.generate(
                        texts,
                        style,
                        latent_hw=(latent_h, latent_w),
                        num_steps=self.num_steps,
                        cfg_scale=self.cfg_scale,
                    )
                for k, img in zip(chunk, imgs, strict=True):
                    inks[k] = _ink_from_pil(ink_crop(to_pil(img)), jobs[k].height)
        return [ink for ink in inks if ink is not None]  # index order preserved; every slot is filled


def render_boxes(boxes: list[LineBox], renderer: LineRenderer) -> list[np.ndarray]:
    """Render each box to ink — the per-line path. Printed form-furniture (table headers, row numbers) is
    typeset; handwriting goes to ``renderer``. The batched equivalent is :func:`render_layouts`."""
    typeset = PlaceholderRenderer()
    return [
        typeset.render(box.text, box.h) if box.printed else renderer.render(box.text, box.h) for box in boxes
    ]


def render_layouts(
    renderer: BatchedDiffuRenderer, layouts: list[Layout], box_styles: list[list[str]]
) -> list[list[np.ndarray]]:
    """Batch-render every HANDWRITING line of every layout in one model pass; typeset (``printed``) boxes
    — form headers, row numbers — are rendered separately. ``box_styles[li][bi]`` is the writer-style path
    for that box, so a page can mix hands (e.g. a different style for marginalia). Returns inks grouped per
    layout (1:1 boxes)."""
    typeset = PlaceholderRenderer()
    jobs: list[LineJob] = []
    job_pos: list[tuple[int, int]] = []
    printed: dict[tuple[int, int], np.ndarray] = {}
    for li, (layout, styles) in enumerate(zip(layouts, box_styles, strict=True)):
        for bi, box in enumerate(layout.boxes):
            if box.printed:
                printed[(li, bi)] = typeset.render(box.text, box.h)
            else:
                jobs.append(LineJob(text=box.text, height=box.h, style_path=styles[bi]))
                job_pos.append((li, bi))
    model_inks = renderer.render_jobs(jobs)
    by_pos = dict(zip(job_pos, model_inks, strict=True))
    return [
        [printed[(li, bi)] if (li, bi) in printed else by_pos[(li, bi)] for bi in range(len(layout.boxes))]
        for li, layout in enumerate(layouts)
    ]