Upload pxg_tiny/pipeline.py with huggingface_hub

pxg_tiny/pipeline.py

@@ -0,0 +1,129 @@
+"""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)