pxg_tiny/pipeline.py

"""PXG-Tiny high-level pipeline (offline, NumPy-only).

PXGPipeline wraps OfflinePipeline with:
  * generate_pixels(text)     -> uint8 grid (16x16 palette indices)
  * generate_png(text, path)  -> scaled RGBA PNG
  * generate_turnaround(text, seeds) / variations
  * should_ask / ask          -> rule-based clarify/refuse gate (ask-first)
  * auto-retry sampling scored by quality.check_sprite (reject-and-log)
"""
import sys
from collections import Counter
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from pxg_tiny.runtime_pipeline import OfflinePipeline  # noqa: E402
from pxg_tiny import quality as Q  # noqa: E402
from pxg_tiny.render import save_png  # noqa: E402

DEFAULT_BUNDLE = Path(__file__).resolve().parents[2] / "weights"


class PXGPipeline:
    def __init__(self, bundle_dir=None):
        self.pipe = OfflinePipeline(bundle_dir or DEFAULT_BUNDLE)

    # ------------------------------------------------------------- gate --
    @staticmethod
    def should_ask(text):
        return Q.should_ask(text)

    @classmethod
    def ask(cls, text):
        """Return a clarifying question / refusal message, or None if the
        prompt is acceptable."""
        label, msg = cls.should_ask(text)
        return msg if label in ("clarify", "refuse") else None

    # -------------------------------------------------------- sampling --
    @staticmethod
    def _anchor_ids(spec):
        """Canonical corpus-style caption ids for the anchor retry stage —
        prompt normalization using only our own parser (offline)."""
        from pxg_tiny.config import encode_caption
        cls = spec.get("cls")
        cap = Q.ANCHOR_CAPTIONS.get(cls)
        if cap is None:
            return None
        if (spec.get("material") and cls not in ("wizard", "archer", "zombie")
                and f" {spec['material']}" not in cap):
            art = "an" if spec["material"][0] in "aeiou" else "a"
            cap = cap.replace("a ", f"{art} {spec['material']} ", 1)
        return np.array(encode_caption(cap), dtype=np.int64)

    def generate_pixels(self, text, seed=0, temperature=None, top_k=None,
                        retries=8, enforce_quality=True):
        """English instruction -> 16x16 grid of palette indices.

        Escalating verify-and-retry schedule (all offline, self-contained):
          attempt 0     raw sample
          attempts 1-3  self-guidance: palette-logit bias, strength ramps
          attempts 4-5  + spatial bias (face box / vial margins)
          attempts 6-7  + canonical anchor caption (prompt normalization)
        The first sprite passing the grounding checks wins; reject reasons
        are counted and returned in the metadata dict."""
        label, msg = self.should_ask(text)
        if label != "accept":
            return None, {"gate": label, "message": msg}

        spec = Q.parse_prompt(text)
        anchor_ids = self._anchor_ids(spec)
        struct_prefix = Q.STRUCTURAL_PREFIX.get(spec.get("cls"))
        rejects = Counter()
        grid = None
        for k in range(max(1, retries)):
            lb = None
            ids = None
            ptoks = struct_prefix if (struct_prefix is not None and k >= 3) else None
            if k >= 1:
                lb = Q.bias_from_spec(spec, strength=1.6 + 0.5 * min(k, 5))
            if k >= 4:
                combined = np.zeros((256, 32), dtype=np.float64)
                if lb is not None:
                    combined[:] = np.asarray(lb, dtype=np.float64)[None, :]
                combined += Q.positional_bias_from_spec(
                    spec, strength=1.4 + 0.3 * (k - 4))
                lb = combined
            if k >= 6 and anchor_ids is not None:
                ids = anchor_ids                       # canonical anchor
            grid = self.pipe.generate_grid(text, seed=seed + 1013 * k,
                                           temperature=temperature,
                                           top_k=top_k, logit_bias=lb,
                                           ids_override=ids,
                                           prefix_tokens=ptoks)
            if enforce_quality:
                ok, reasons = Q.check_sprite(grid, spec)
                if ok:
                    return grid, {"gate": "accept", "seed": seed + 1013 * k,
                                  "attempt": k, "rejects": dict(rejects)}
                for r in reasons:
                    rejects[r] += 1
            else:
                return grid, {"gate": "accept", "seed": seed, "attempt": k,
                              "rejects": {}}
        # fall back to the last sample rather than failing hard
        return grid, {"gate": "accept_degraded", "seed": seed,
                      "attempt": retries, "rejects": dict(rejects)}

    def generate_png(self, text, path, scale=8, seed=0, **kw):
        grid, meta = self.generate_pixels(text, seed=seed, **kw)
        if grid is None:
            return None, meta
        save_png(grid, Path(path), scale=scale)
        meta["path"] = str(path)
        return grid, meta

    def generate_turnaround(self, text, seeds=(1, 2, 3, 4), **kw):
        """Several consistent samples of the same instruction (seed family);
        tiny-model analog of the big sibling's multiview turnaround."""
        grids = []
        for s in seeds:
            g, m = self.generate_pixels(text, seed=s, **kw)
            grids.append((g, m))
        return grids

    def variations(self, text, k=4, start_seed=0, **kw):
        return self.generate_turnaround(text, seeds=tuple(
            start_seed + 37 * i for i in range(k)), **kw)