Spaces:
Running
Running
| """Browser-side generation driver for PXG-Tiny (Pyodide, no PIL). | |
| Mirrors the shipped PXGPipeline escalation schedule exactly, but without the | |
| render/PIL dependency: | |
| attempt 0 raw sample | |
| attempts 1-3 palette-logit self-guidance, strength ramps | |
| attempts 4-5 + spatial priors (face box / vial margins) | |
| attempts 6-7 + canonical anchor caption | |
| Structural prefixes (deer antlers) engage from attempt 3, as shipped. | |
| """ | |
| import numpy as np | |
| from pxg_tiny import quality as Q | |
| from pxg_tiny.config import encode_caption | |
| from pxg_tiny.runtime_pipeline import OfflinePipeline | |
| class WebPipeline: | |
| def __init__(self, bundle_dir): | |
| self.pipe = OfflinePipeline(bundle_dir) | |
| def _anchor_ids(self, spec): | |
| 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(self, text, seed=0, retries=8, enforce_quality=True): | |
| label, msg = Q.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 = {} | |
| 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 and enforce_quality: | |
| lb = Q.bias_from_spec(spec, strength=1.6 + 0.5 * min(k, 5)) | |
| if k >= 4 and enforce_quality: | |
| 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 | |
| grid = self.pipe.generate_grid(text, seed=seed + 1013 * k, | |
| logit_bias=lb, | |
| ids_override=ids, | |
| prefix_tokens=ptoks) | |
| if not enforce_quality: | |
| return grid, {"gate": "accept", "seed": seed, "attempt": k, | |
| "rejects": {}} | |
| ok, reasons = Q.check_sprite(grid, spec) | |
| if ok: | |
| return grid, {"gate": "accept", "seed": seed + 1013 * k, | |
| "attempt": k, "rejects": rejects} | |
| for r in reasons: | |
| rejects[r] = rejects.get(r, 0) + 1 | |
| return grid, {"gate": "accept_degraded", "seed": seed, | |
| "attempt": retries, "rejects": rejects} | |
| def grid_rgba(self, grid): | |
| rgba = np.zeros((16, 16, 4), dtype=np.uint8) | |
| for idx in range(1, len(self.pipe.palette)): | |
| m = grid == idx | |
| rgba[m] = self.pipe.palette[idx] | |
| return rgba | |