Blame for reference.py

"""Oasis reference: Mirage's forward pass in PyTorch, the 2-bit quantiser, the text reader's
quantisation-aware training and evaluation, and the export to .oasis.

    python reference.py COMMAND [options]        (python reference.py COMMAND --help describes one)

Model     export, checkpoint, run, det-eval
Reader    train, eval-reader, eval-detected, pack-words, pack-detected, synth-val

The forward pass follows the equations Cactus Compute published for Mirage, and its decoder follows the
Needle 3 reference in cactus-compute (Apache-2.0). Commands read Mirage's mirage.safetensors and mirage.cact
from mirage/ beside this file (--mirage), and data from data/ (--data): COCO val2017 photos in
data/coco/val300 (python oasis.py fetch-coco) with data/coco/annotations/instances_val2017.json, and word
crops in parquet files. The engine's benchmarks against Mirage are in oasis.py.
"""
import argparse
import glob
import io
import json
import math
import os
import random
import re
import struct
import sys
import time
import unicodedata
import zlib
from dataclasses import dataclass
from math import prod

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image, ImageDraw, ImageFilter, ImageFont
from safetensors.torch import load_file, save_file

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
from oasis import OasisEngine, coco_truth, edit_distance, iou, sha256, val300, word_key, word_scene  # noqa: E402

CHARSET_ENTRY, VOCAB_ENTRY = 804, 805   # mirage.cact entries holding the reader's charset and the tokenizer
EMBEDDING_ROWS = 8299
LINUX_FONTS = ["/usr/share/fonts/truetype/dejavu/*.ttf", "/usr/share/fonts/truetype/liberation/*.ttf"]



# ===================================================================================== shared helpers
def parser(name, doc):
    """A command's argument parser, with the --mirage and --data directories every command accepts."""
    ap = argparse.ArgumentParser(prog=f"reference.py {name}", description=doc,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--mirage", default=os.path.join(HERE, "mirage"), help="directory of mirage.safetensors and mirage.cact")
    ap.add_argument("--data", default=os.path.join(HERE, "data"), help="directory of the benchmark data")
    return ap


def mirage_weights(args):
    return os.path.join(args.mirage, "mirage.safetensors")


def mirage_archive(args):
    return os.path.join(args.mirage, "mirage.cact")


def charset_of(arc):
    return "".join(chr(int(c)) for c in arc.array(CHARSET_ENTRY))


def vocab_of(arc, rows=EMBEDDING_ROWS):
    return Vocab(arc.raw(VOCAB_ENTRY), embedding_rows=rows)


def dead_classes(weights):
    """Reader classes whose output weights are zero in `weights`: characters Mirage's reader never learned."""
    k = weights["reader/out/kernel"]
    return [c for c in range(1, k.shape[1]) if k[:, c].norm() < 1e-6]


def vocab_words(vocab):
    """Alphabetic pieces of at least two letters, the word list of the synthetic text."""
    return sorted({p.replace("▁", "") for p, t in zip(vocab.pieces, vocab.types)
                   if t == 0 and len(p.replace("▁", "")) >= 2 and p.replace("▁", "").isalpha()})


def ctc_text(labels, charset):
    """Collapse a CTC label path: drop blanks (0) and repeats."""
    s, prev = [], 0
    for lab in labels:
        if lab != 0 and lab != prev:
            s.append(charset[lab - 1])
        prev = lab
    return "".join(s)


def cq_plan():
    """Mirage's quantised matrices in mirage.cact order: the archive entry, the float32 tensor, its decoder
    layer for tensors stacked over the eight layers, how its [out, in] matrix is read (T: transposed;
    T_per_layer: transposed per layer and stacked; reshape: as stored), the bits and the [out, in] shape."""
    plan = []

    def add(entry, name, shape, how="T", bits=2, layer=None):
        plan.append(dict(entry=entry, name=name, layer=layer, how=how, bits=bits, shape=list(shape)))

    add(0, "embedding/embedding", (8299, 512), "reshape", 4)
    for l in range(8):
        for k, proj, rows in ((0, "q", 384), (1, "k", 96), (2, "v", 128), (8, "gate", 512), (9, "out", 512)):
            add(2 + 27 * l + k, f"stack/layers/block/self_attn/{proj}_proj/kernel", (rows, 512), layer=l)
    for k, part, rows in ((0, "pre", 32), (1, "post", 32), (2, "res", 128)):
        add(223 + k, f"stack/mhc_phi_{part}", (rows, 2048), "T_per_layer", 4)
    for i in range(2):
        add(228 + 4 * i, f"engrams_{i}/embedding", (73728, 128), "reshape")
        add(229 + 4 * i, f"engrams_{i}/key_proj/kernel", (512, 512))
        add(230 + 4 * i, f"engrams_{i}/value_proj/kernel", (512, 512))
    for l in range(8):
        for k, proj, rows in ((0, "q", 384), (1, "k", 384), (2, "v", 512), (3, "gate", 512), (4, "out", 512)):
            add(239 + 10 * l + k, f"stack/layers/block/cross_attn/{proj}_proj/kernel", (rows, 512), layer=l)
    for t in range(12):
        for k, proj, shape in ((0, "q", (128, 512)), (1, "k", (32, 512)), (2, "v", (64, 512)), (3, "gate", (256, 512)),
                               (4, "out", (512, 256))):
            add(542 + 15 * t + k, f"top_{t}/self_attn/{proj}_proj/kernel", shape)
    add(773, "slot_pool/kernel", (512, 512))
    for j, shape in enumerate(((64, 32), (128, 64), (384, 128))):
        add(777 + 3 * j, f"reader/pw_{j}/kernel", shape)
    for b in range(4):
        add(787 + 4 * b, f"reader/up_{b}/kernel", (768, 384))
        add(788 + 4 * b, f"reader/down_{b}/kernel", (384, 384))
    add(802, "reader/out/kernel", (307, 384))
    return plan


# ===================================================================================== tokenizer
# SentencePiece vocabulary carried in a .cact RAW entry, plus Mirage's added language and reference tokens.
# RAW layout: u32 piece count, u32 pad, eos, bos, unk ids, u8 add-dummy-prefix flag, u8 byte-fallback flag,
# u16 padding, then per piece f32 score, u8 type, u16 length, UTF-8 bytes. Types: 0 normal, 1 unknown,
# 2 control, 3 user-defined, 4 byte. Mirage's embedding has 96 rows past the 8,203 pieces: reference tokens
# for 64 object regions and 32 text regions.
N_OBJECT_REFS = 64
N_TEXT_REFS = 32


class Vocab:
    def __init__(self, raw: bytes, embedding_rows: int):
        count, self.pad, self.eos, self.bos, self.unk = struct.unpack_from("<5I", raw, 0)
        self.add_dummy_prefix, self.byte_fallback = bool(raw[20]), bool(raw[21])
        off = 24
        self.pieces, self.types, self.scores = [], [], []
        for _ in range(count):
            score, typ, ln = struct.unpack_from("<fBH", raw, off)
            off += 7
            self.pieces.append(raw[off:off + ln].decode("utf-8", "replace"))
            self.types.append(typ)
            self.scores.append(score)
            off += ln
        self.n_pieces = count
        self.ref_base = count
        self.size = embedding_rows
        self.id_of = {p: i for i, p in enumerate(self.pieces)}
        self.lang_ids = {p[2:-2]: i for i, p in enumerate(self.pieces) if re.fullmatch(r"<\|[a-z]{2}\|>", p)}

    def ref_id(self, kind, k):
        return self.ref_base + (k if kind == "object" else N_OBJECT_REFS + k)

    def piece(self, i):
        if i >= self.n_pieces:
            k = i - self.ref_base
            return f"<obj{k}>" if k < N_OBJECT_REFS else f"<txt{k - N_OBJECT_REFS}>"
        return self.pieces[i]

    def decode(self, ids, show_special=True):
        out, buf = [], bytearray()

        def flush():
            if buf:
                out.append(buf.decode("utf-8", "replace"))
                buf.clear()

        for i in ids:
            if i >= self.n_pieces:
                flush()
                out.append(self.piece(i))
                continue
            t, p = self.types[i], self.pieces[i]
            if t == 4:
                m = re.fullmatch(r"<0x([0-9A-Fa-f]{2})>", p)
                if m:
                    buf.append(int(m.group(1), 16))
                    continue
            flush()
            if t in (1, 2, 3):
                if show_special:
                    out.append(p)
            else:
                out.append(p.replace("▁", " "))
        flush()
        return "".join(out)


# ===================================================================================== Cactus archive
# Cactus .cact archives (tag 0x05E12A84): a 196-byte header of 49 little-endian 32-bit fields (the last is
# rope_theta as float32), codebook floats, then a 44-byte record per tensor: u8 dtype, u8 ndim, u16 pad,
# u32 shape[4], u64 offset, u64 nbytes, u32 group, u32 bits. Dtypes: 1 fp16, 2 fp32, 3 CQ, 4 RAW.
# A CQ tensor is stored [out, in]; each row is split into groups of `group` weights (zero-padded), every
# group is rotated by the natural-order normalised Walsh-Hadamard matrix, its L2 norm is kept in fp16, and
# each rotated coordinate of the unit direction is replaced by the index of the nearest Lloyd-Max codebook
# entry. The blob holds every row's packed indices (LSB first), then every group norm.
TAG_NEEDLE3 = 0x05E12A84
DTYPES = {1: "fp16", 2: "fp32", 3: "cq", 4: "raw"}
CODEBOOK_SLICES = {2: slice(0, 4), 3: slice(4, 12), 4: slice(12, 28)}


@dataclass
class Entry:
    index: int
    dtype: str
    shape: tuple
    offset: int
    nbytes: int
    group: int
    bits: int


def hadamard(n):
    h = np.array([[1.0]])
    while h.shape[0] < n:
        h = np.block([[h, h], [h, -h]])
    return h / np.sqrt(n)


class Archive:
    def __init__(self, path):
        with open(path, "rb") as f:
            self.data = f.read()
        b = self.data
        self.fields = struct.unpack_from("<48I", b, 0)
        self.rope_theta = struct.unpack_from("<f", b, 192)[0]
        self.tag, n_tensors, cb_len = self.fields[0], self.fields[1], self.fields[2]
        if self.tag != TAG_NEEDLE3:
            raise ValueError(f"unsupported archive tag 0x{self.tag:08X}")
        self.codebook = np.array(struct.unpack_from(f"<{cb_len}f", b, 196), dtype=np.float64)
        off = 196 + 4 * cb_len
        self.entries = []
        for i in range(n_tensors):
            dtype, ndim = b[off], b[off + 1]
            shape = struct.unpack_from("<4I", b, off + 4)[:ndim]
            start, nbytes = struct.unpack_from("<QQ", b, off + 20)
            group, bits = struct.unpack_from("<II", b, off + 36)
            self.entries.append(Entry(i, DTYPES[dtype], tuple(shape), start, nbytes, group, bits))
            off += 44

    def raw(self, i):
        e = self.entries[i]
        return self.data[e.offset:e.offset + e.nbytes]

    def array(self, i):
        """Return entry i as float64 in its stored layout (CQ dequantised)."""
        e = self.entries[i]
        if e.dtype in ("fp16", "fp32"):
            dt = np.float16 if e.dtype == "fp16" else np.float32
            a = np.frombuffer(self.data, dtype=dt, count=e.nbytes // np.dtype(dt).itemsize, offset=e.offset)
            return a.astype(np.float64).reshape(e.shape)
        if e.dtype == "cq":
            return self._dequant(e)
        raise TypeError("RAW entries have no array form")

    def cq_parts(self, e):
        rows, cols = prod(e.shape[:-1]), e.shape[-1]
        ng = -(-cols // e.group)
        padded = ng * e.group
        blob = np.frombuffer(self.data, dtype=np.uint8, count=e.nbytes, offset=e.offset)
        ib = rows * padded * e.bits // 8
        per = 8 // e.bits
        idx = np.empty(ib * per, dtype=np.uint8)
        for k in range(per):
            idx[k::per] = (blob[:ib] >> (e.bits * k)) & ((1 << e.bits) - 1)
        idx = idx.reshape(rows, ng, e.group)
        norms = blob[ib:].view(np.float16).astype(np.float64).reshape(rows, ng)
        return idx, norms, rows, cols, padded

    def _dequant(self, e):
        idx, norms, rows, cols, padded = self.cq_parts(e)
        cb = self.codebook[CODEBOOK_SLICES[e.bits]]
        rot = cb[idx] * norms[..., None]
        w = rot @ hadamard(e.group)
        return w.reshape(rows, padded)[:, :cols].reshape(e.shape)


# ===================================================================================== quantisation
# A quantised matrix is held [out, in]; each row is cut into 128-weight groups, every group is rotated by
# the normalised natural-order Walsh-Hadamard matrix, and each rotated coordinate is replaced by a codebook
# entry times one fp16 scale per group. "stock" sets the scale to the group's L2 norm and picks the nearest
# entry of the unit-norm direction (the published scheme); "refit" alternates nearest-entry assignment and
# the least-squares scale.
H128 = torch.tensor(hadamard(128), dtype=torch.float64)


def nearest(x, cb):
    """Index of the nearest codebook entry for every element of x (codebook sorted ascending)."""
    mids = (cb[1:] + cb[:-1]) / 2
    return torch.bucketize(x.contiguous(), mids)


def quantize_rows(W, cb, method="stock", iters=8, chunk=8192):
    rows = W.shape[0]
    if rows > chunk:
        return torch.cat([quantize_rows(W[i:i + chunk], cb, method, iters, chunk) for i in range(0, rows, chunk)])
    rows, cols = W.shape
    pc = -(-cols // 128) * 128
    Wp = torch.zeros(rows, pc, dtype=torch.float64)
    Wp[:, :cols] = W.double()
    g = Wp.view(rows, pc // 128, 128) @ H128
    n = g.norm(dim=-1, keepdim=True)
    idx = nearest(g / n.clamp_min(1e-30), cb)
    s = n
    if method == "refit":
        for _ in range(iters):
            c = cb[idx]
            s = (g * c).sum(-1, keepdim=True) / (c * c).sum(-1, keepdim=True).clamp_min(1e-30)
            safe = torch.where(s.abs() > 1e-30, s, torch.full_like(s, 1e-30))
            idx = nearest(g / safe, cb)
        c = cb[idx]
        s = (g * c).sum(-1, keepdim=True) / (c * c).sum(-1, keepdim=True).clamp_min(1e-30)
    s = s.half().double()
    rec = (cb[idx] * s) @ H128
    return rec.reshape(rows, pc)[:, :cols].float()


def plan_get(p, item):
    """The [out, in] matrix of a plan item, read from the weight dictionary p."""
    a = p[item["name"]]
    if item["layer"] is not None:
        a = a[item["layer"]]
    if item["how"] == "T":
        return a.T
    if item["how"] == "T_per_layer":
        L, i_, o_ = a.shape
        return a.transpose(1, 2).reshape(L * o_, i_)
    return a.reshape(item["shape"])


def plan_put(p, item, W):
    """Write the [out, in] matrix W of a plan item back into the weight dictionary p."""
    a = p[item["name"]]
    if item["how"] == "T":
        new = W.T.contiguous()
    elif item["how"] == "T_per_layer":
        L, i_, o_ = a.shape
        new = W.reshape(L, o_, i_).transpose(1, 2).contiguous()
    else:
        new = W.reshape(a.shape[1:] if item["layer"] is not None else a.shape)
    if item["layer"] is not None:
        a[item["layer"]] = new
    else:
        p[item["name"]] = new


def variant(p32, plan, codebook, method="stock", bits_rules=()):
    """Return a copy of p32 with every planned tensor quantised and everything else rounded to fp16.

    bits_rules: sequence of (regex on tensor name, bits) applied in order to override the plan.
    """
    p = {k: v.half().float().clone() for k, v in p32.items()}
    size_bits = 0
    for item in plan:
        bits = item["bits"]
        for pat, b in bits_rules:
            if re.search(pat, item["name"]):
                bits = b
        cb = torch.tensor(codebook[CODEBOOK_SLICES[bits]], dtype=torch.float64)
        W = plan_get(p32, item)
        plan_put(p, item, quantize_rows(W, cb, method))
        rows, cols = W.shape
        size_bits += rows * (-(-cols // 128) * 128) * bits + rows * (-(-cols // 128)) * 16
    quant_names = {it["name"] for it in plan}
    fp16_params = sum(v.numel() for k, v in p32.items() if k not in quant_names and not k.startswith(("mtp_", "name_head")))
    return p, dict(cq_mb=round(size_bits / 8 / 1e6, 3), fp16_mb=round(fp16_params * 2 / 1e6, 3))


class Quantizer:
    """Quantisation-aware training for Hadamard-codebook weights: fake_quant_kernel(W) applies the stock
    quantiser in the forward pass (each 128-weight input group of a row rotated by the Walsh-Hadamard matrix,
    scaled by its fp16 L2 norm, every coordinate snapped to the nearest codebook entry) and passes gradients
    straight through to the latent weight."""

    def __init__(self, codebook, device):
        self.H = torch.tensor(hadamard(128), dtype=torch.float32, device=device)
        self.cb = {b: torch.tensor(codebook[CODEBOOK_SLICES[b]], dtype=torch.float32, device=device) for b in (2, 3, 4)}
        self.mid = {b: (c[1:] + c[:-1]) / 2 for b, c in self.cb.items()}

    def quantize_rows(self, R, bits):
        """R: [rows, cols] in [out, in] layout. Returns the dequantised quantised matrix."""
        rows, cols = R.shape
        pc = -(-cols // 128) * 128
        if pc != cols:
            R = torch.nn.functional.pad(R, (0, pc - cols))
        g = R.reshape(rows, pc // 128, 128) @ self.H
        n = g.norm(dim=-1, keepdim=True)
        idx = torch.bucketize((g / n.clamp_min(1e-30)).contiguous(), self.mid[bits])
        s = n.half().float()
        rec = (self.cb[bits][idx] * s) @ self.H
        return rec.reshape(rows, pc)[:, :cols]

    def fake_quant_kernel(self, W, bits):
        """W: a Dense kernel stored [in, out]; quantised along rows of W.T."""
        with torch.no_grad():
            Q = self.quantize_rows(W.T.float(), bits).T
        return W + (Q - W).detach()

# ===================================================================================== forward pass
# Mirage's forward pass, float32 on the CPU, from the weights of mirage.safetensors.
EPS = 1e-6
D_MODEL = 512
CANVAS = 512
N_OBJ, N_TXT = 64, 32
ENGRAM_SITES = (3, 7)
ENGRAM_ORDERS = (2, 3)
ENGRAM_HEADS = 2
ENGRAM_SLOTS = 18432
ENGRAM_SEED = 0x9E3779B9
ENGRAM_PRIME = 0x01000193
ENGRAM_TAPS, ENGRAM_DILATION = 4, 3

# Settings the published equations leave open; kept together so they can be varied.
OPTS = {
    "pad_mode": "sym",        # "sym": pad 1 on both sides; "same": TF-style (0, 1) for stride 2
    "canvas_fill": 0.0,       # normalised value of letterbox padding
    "rope2d_rows_first": True,
    "rope2d_layout": "pairs16",  # "pairs16": pairs (r, r+16), rows on r<8; "blocks": two 16-d halves
    "walsh": "sequency",        # or "natural"
    "level_gate": "raw",        # or "sigmoid"
    "resample": Image.BILINEAR,
}


def zc_norm(x, scale, dim=-1):
    shape = [1] * x.dim()
    shape[dim] = -1
    return x * torch.rsqrt(x.pow(2).mean(dim, keepdim=True) + EPS) * (1 + scale.view(shape))


def rms_unit(x):
    return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + EPS)


def _pad(x, stride_hw):
    if OPTS["pad_mode"] == "sym" or stride_hw == (1, 1):
        return F.pad(x, (1, 1, 1, 1))
    pads = []
    for size, s in ((x.shape[3], stride_hw[1]), (x.shape[2], stride_hw[0])):
        out = -(-size // s)
        total = max((out - 1) * s + 3 - size, 0)
        pads += [total // 2, total - total // 2]
    return F.pad(x, pads)


def conv3(x, w_hwio, stride=(1, 1)):
    w = w_hwio.permute(3, 2, 0, 1).contiguous()
    return F.conv2d(_pad(x, stride), w, stride=stride)


def dwconv(x, w_hw1c, stride=(1, 1)):
    w = w_hw1c.permute(3, 2, 0, 1).contiguous()
    if w.shape[-1] == 3 and w.shape[-2] == 3:
        return F.conv2d(_pad(x, stride), w, stride=stride, groups=x.shape[1])
    kh, kw = w.shape[-2:]
    return F.conv2d(F.pad(x, (kw // 2, kw // 2, kh // 2, kh // 2)), w, stride=stride, groups=x.shape[1])


def pw(x, k):
    return (x.permute(0, 2, 3, 1) @ k).permute(0, 3, 1, 2)


def cnorm(x, scale):
    return zc_norm(x, scale, dim=1)


def walsh_sequency(n):
    q = n.bit_length() - 1
    j = np.arange(n)
    g = j ^ (j >> 1)
    s = np.array([int(format(int(v), f"0{q}b")[::-1], 2) for v in g])
    pc = np.vectorize(lambda v: bin(int(v)).count("1"))(s[:, None] & np.arange(n)[None, :])
    return torch.tensor(((-1.0) ** pc) / math.sqrt(n), dtype=torch.float32)


def band(u, w):
    r = (w.shape[1] - 1) // 2
    win = F.pad(u, (r, r)).unfold(-1, 2 * r + 1, 1)
    return (win * w).sum(-1)


def sinkhorn(logits, iters=20):
    k = logits
    for _ in range(iters):
        k = k - torch.logsumexp(k, dim=-1, keepdim=True)
        k = k - torch.logsumexp(k, dim=-2, keepdim=True)
    return k.exp()


def rope_half(x, pos, theta):
    hd = x.shape[-1]
    freqs = 1.0 / (theta ** (torch.arange(0, hd, 2, dtype=torch.float32) / hd))
    ang = pos.float()[..., None] * freqs
    c, s = ang.cos(), ang.sin()
    x1, x2 = x[..., : hd // 2], x[..., hd // 2:]
    return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s], -1)


def sincos_2d(rows, cols, d=D_MODEL, theta=10000.0):
    def one(pos, dd):
        omega = 1.0 / (theta ** (torch.arange(dd // 2, dtype=torch.float32) / (dd // 2)))
        a = pos.float()[:, None] * omega[None]
        return torch.cat([a.sin(), a.cos()], -1)
    return torch.cat([one(rows, d // 2), one(cols, d // 2)], -1)


def engram_index(tokens, t, oi, order, h):
    acc = (ENGRAM_SEED * (oi * ENGRAM_HEADS + h + 1)) & 0xFFFFFFFF
    for j in range(order):
        tok = tokens[t - j] if t - j >= 0 else 0
        acc = ((acc ^ tok) * ENGRAM_PRIME) & 0xFFFFFFFF
    acc ^= acc >> 15
    return acc % ENGRAM_SLOTS


class Mirage:
    def __init__(self, weights_path):
        self.p = {k: v.float() for k, v in load_file(weights_path).items()}
        self.W = walsh_sequency(D_MODEL)
        self.W_nat = torch.tensor(hadamard(D_MODEL), dtype=torch.float32)
        p = self.p
        self.perm1 = torch.tensor(np.random.RandomState(11).permutation(D_MODEL))
        self.perm2 = torch.tensor(np.random.RandomState(13).permutation(D_MODEL))
        self.T = p["embedding/embedding"]

    # ---------------------------------------------------------------- image
    def letterbox(self, img):
        img = img.convert("RGB")
        W0, H0 = img.size
        r = CANVAS / max(W0, H0)
        w, h = max(1, round(W0 * r)), max(1, round(H0 * r))
        rs = np.asarray(img.resize((w, h), OPTS["resample"]), dtype=np.float32) / 127.5 - 1.0
        canvas = np.full((CANVAS, CANVAS, 3), OPTS["canvas_fill"], dtype=np.float32)
        canvas[:h, :w] = rs
        x = torch.from_numpy(canvas).permute(2, 0, 1)[None]
        return x, dict(W=W0, H=H0, r=r, w=w, h=h)

    @staticmethod
    def valid_mask(meta, stride):
        n = CANVAS // stride
        m = torch.zeros(1, 1, n, n)
        m[..., : -(-meta["h"] // stride), : -(-meta["w"] // stride)] = 1.0
        return m

    def hconvglu(self, x, pre):
        p = self.p
        W = self.W_nat if OPTS["walsh"] == "natural" else self.W
        u = x @ W
        g = F.gelu(band(u, p[pre + "/gate_conv"]) + p[pre + "/gate_bias"], approximate="tanh")
        v = band(u, p[pre + "/value_conv"])
        return band(g * v, p[pre + "/out_conv"]) @ W

    # --------------------------------------------------------------- encoder
    @torch.inference_mode()
    def encode(self, img):
        p = self.p
        x, meta = self.letterbox(img)
        ms = {s: self.valid_mask(meta, s) for s in (2, 4, 8, 16, 32)}
        f = ms[2] * cnorm(conv3(x, p["stem/w0"], (2, 2)), p["stem/norm_0/scale"])
        feats = {2: f}
        for j in range(1, 5):
            s = 2 ** (j + 1)
            y = dwconv(F.silu(f), p[f"stem/dw_{j}"], (2, 2))
            f = ms[s] * cnorm(pw(y, p[f"stem/pw_{j}/kernel"]), p[f"stem/norm_{j}/scale"])
            feats[s] = f
        S4, S8, S16 = F.silu(feats[4]), F.silu(feats[8]), F.silu(feats[16])
        X = F.silu(feats[32])
        m32 = ms[32]
        for i in range(16):
            X = self.body_block(X, i, m32)
        Z = X.flatten(2).transpose(1, 2)                       # [1, 256, 512]
        Z = torch.cat([Z, p["registers"][None]], 1)            # [1, 260, 512]
        valid = torch.cat([m32.flatten().bool(), torch.ones(4, dtype=torch.bool)])
        rows = torch.arange(256) // 16
        cols = torch.arange(256) % 16
        levels = [Z]
        for l in range(12):
            Z = self.top_block(Z, l, valid, rows, cols)
            levels.append(Z)
        mv = m32.flatten()[None, :, None]
        E = zc_norm(levels[12][:, :256], p["enc_norm/scale"])
        for l in range(12):
            g = p["level_gates"][l] if OPTS["level_gate"] == "raw" else torch.sigmoid(p["level_gates"][l])
            E = E + g * zc_norm(levels[l][:, :256], p[f"level_norm_{l}/scale"])
        E = mv * E                                             # [1, 256, 512]
        D = self.neck(E, S16, S8, ms)
        return dict(meta=meta, ms=ms, E=E, D=D, S4=S4, valid=m32.flatten().bool())

    def body_block(self, X, i, m):
        p, pre = self.p, f"body_{i}"
        Z = cnorm(X, p[pre + "/mix_norm/scale"])
        v = dwconv(Z, p[pre + "/mix/dw_value"]) + p[pre + "/mix/b_value"].view(1, -1, 1, 1)
        g = dwconv(Z, p[pre + "/mix/dw_gate"]) + p[pre + "/mix/b_gate"].view(1, -1, 1, 1)
        Xt = X + m * dwconv(v * F.silu(g), p[pre + "/mix/dw_down"])
        xt = Xt.permute(0, 2, 3, 1)
        H = self.hconvglu(zc_norm(xt, p[pre + "/hconv_norm/scale"]), pre + "/hconv")
        mm = m.permute(0, 2, 3, 1)
        r = torch.sqrt(((mm * H) ** 2).sum((1, 2), keepdim=True) + EPS)
        G = H + p[pre + "/grn/gamma"] * H * r / (r.mean(-1, keepdim=True) + EPS) + p[pre + "/grn/beta"]
        return (xt + mm * G).permute(0, 3, 1, 2)

    def top_block(self, Z, l, valid, rows, cols):
        p, pre = self.p, f"top_{l}/"
        T = Z.shape[1]
        z = zc_norm(Z, p[pre + "attn_norm/scale"])
        q = (z @ p[pre + "self_attn/q_proj/kernel"]).view(1, T, 4, 32).transpose(1, 2)
        k = (z @ p[pre + "self_attn/k_proj/kernel"]).view(1, T, 1, 32).transpose(1, 2)
        v = (z @ p[pre + "self_attn/v_proj/kernel"]).view(1, T, 1, 64).transpose(1, 2)
        q = zc_norm(q, p[pre + "self_attn/q_norm/scale"])
        k = zc_norm(k, p[pre + "self_attn/k_norm/scale"])
        a, b = (rows, cols) if OPTS["rope2d_rows_first"] else (cols, rows)

        def rope2(t):
            c = t[:, :, :256]
            if OPTS["rope2d_layout"] == "blocks":
                c = torch.cat([rope_half(c[..., :16], a, 10000.0), rope_half(c[..., 16:], b, 10000.0)], -1)
            else:
                freqs = 1.0 / (10000.0 ** (torch.arange(0, 16, 2, dtype=torch.float32) / 16))
                ang = torch.cat([a.float()[:, None] * freqs, b.float()[:, None] * freqs], -1)
                cs, sn = ang.cos(), ang.sin()
                x1, x2 = c[..., :16], c[..., 16:]
                c = torch.cat([x1 * cs - x2 * sn, x2 * cs + x1 * sn], -1)
            return torch.cat([c, t[:, :, 256:]], 2)

        q, k = rope2(q), rope2(k)
        s = (q @ k.transpose(-1, -2)) / math.sqrt(32)
        s = s.masked_fill(~valid.view(1, 1, 1, T), float("-inf"))
        o = (s.softmax(-1) @ v).transpose(1, 2).reshape(1, T, 256)
        A = (o * torch.sigmoid(z @ p[pre + "self_attn/gate_proj/kernel"])) @ p[pre + "self_attn/out_proj/kernel"]
        mv = valid.view(1, T, 1).float()
        Zt = Z + mv * torch.sigmoid(p[pre + "attn_gate"]) * zc_norm(A, p[pre + "post_attn_norm/scale"])
        return Zt + mv * self.hconvglu(zc_norm(Zt, p[pre + "hconv_norm/scale"]), pre + "hconv")

    def res_block(self, Fm, norm, dw, pwk, m):
        p = self.p
        return m * (Fm + pw(dwconv(F.silu(cnorm(Fm, p[norm])), p[dw]), p[pwk]))

    def neck(self, E, S16, S8, ms):
        p = self.p
        Eg = E.transpose(1, 2).reshape(1, D_MODEL, 16, 16)
        t = F.interpolate(pw(Eg, p["dense/coarse/kernel"]), scale_factor=2, mode="nearest")
        t = ms[16] * (t + pw(S16, p["dense/mid/kernel"]))
        t = self.res_block(t, "dense/neck_norm/scale", "dense/neck_dw", "dense/neck_pw/kernel", ms[16])
        d = ms[8] * (F.interpolate(t, scale_factor=2, mode="nearest") + pw(S8, p["dense/fine/kernel"]))
        for i in range(3):
            d = self.res_block(d, f"dense/block_norm_{i}/scale", f"dense/dw_{i}", f"dense/pw_{i}/kernel", ms[8])
        return ms[8] * F.silu(cnorm(d, p["dense/final_norm/scale"]))   # [1, 128, 64, 64]

    # ------------------------------------------------------------- detection
    @torch.inference_mode()
    def detect(self, enc, conf=0.3, text_conf=0.3, nms_iou=None):
        p = self.p
        out = torch.einsum("bchw,cd->bhwd", enc["D"], p["dense/out/kernel"])[0] + p["dense/bias"]
        m8 = enc["ms"][8][0, 0].bool()
        meta = enc["meta"]
        result = {}
        for kind, ch, limit, thr in (("object", 0, N_OBJ, conf), ("text", 1, N_TXT, text_conf)):
            heat = torch.sigmoid(out[..., ch])
            peak = (heat == F.max_pool2d(heat[None, None], 3, 1, 1)[0, 0]) & m8 & (heat >= thr)
            ii, jj = torch.nonzero(peak, as_tuple=True)
            scores = heat[ii, jj]
            order = torch.argsort(scores, descending=True)
            boxes = []
            for o in order.tolist():
                i, j = int(ii[o]), int(jj[o])
                w = abs(float(out[i, j, 2 + 2 * ch]))
                h = abs(float(out[i, j, 3 + 2 * ch]))
                cx = j + float(out[i, j, 6 + 2 * ch])
                cy = i + float(out[i, j, 7 + 2 * ch])
                b = np.array([cx - w / 2, cy - h / 2, cx + w / 2, cy + h / 2]) / 64.0
                b[[0, 2]] = b[[0, 2]].clip(0, meta["w"] / CANVAS)
                b[[1, 3]] = b[[1, 3]].clip(0, meta["h"] / CANVAS)
                if nms_iou and any(iou(b, x["box"]) > nms_iou for x in boxes):
                    continue
                boxes.append(dict(box=b, score=float(scores[o]), cell=(i, j)))
                if len(boxes) >= limit:
                    break
            result[kind] = boxes
        result["raw"] = out
        return result

    def to_1000(self, b, meta):
        sx = CANVAS / meta["r"] / meta["W"] * 1000
        sy = CANVAS / meta["r"] / meta["H"] * 1000
        return [round(float(b[0] * sx)), round(float(b[1] * sy)), round(float(b[2] * sx)), round(float(b[3] * sy))]

    # ----------------------------------------------------------------- masks
    def head(self, Fm, pre, m):
        p = self.p
        for i in range(2):
            Fm = self.res_block(Fm, f"dense/{pre}_norm_{i}/scale", f"dense/{pre}_dw_{i}", f"dense/{pre}_pw_{i}/kernel", m)
        return pw(F.silu(cnorm(Fm, p[f"dense/{pre}_final_norm/scale"])), p[f"dense/{pre}s/kernel" if pre == "proto" else f"dense/{pre}/kernel"])

    @torch.inference_mode()
    def masks(self, enc, boxes):
        p, ms, D = self.p, enc["ms"], enc["D"]
        U = ms[4] * (F.interpolate(D, scale_factor=2, mode="nearest") + pw(enc["S4"], p["dense/edge/kernel"]))
        P = self.head(U, "proto", ms[4])[0]                      # [32, 128, 128]
        C = self.head(D, "coef", ms[8])[0]                       # [32, 64, 64]
        out = []
        ys = (torch.arange(128) + 0.5) / 128
        for bx in boxes:
            i, j = bx["cell"]
            logit = torch.einsum("r,rhw->hw", C[:, i, j], P)
            b = bx["box"]
            inside = ((ys[:, None] >= b[1]) & (ys[:, None] <= b[3]) & (ys[None, :] >= b[0]) & (ys[None, :] <= b[2]))
            ci, cj = min(127, int((b[1] + b[3]) / 2 * 128)), min(127, int((b[0] + b[2]) / 2 * 128))
            inside[ci, cj] = True
            out.append(torch.sigmoid(logit) * inside)
        return out

    # ---------------------------------------------------------------- reader
    @torch.inference_mode()
    def read(self, img, box_img):
        """box_img: (x0, y0, x1, y1) in original pixels. Returns frame label ids and per-frame probs."""
        img = img.convert("RGB")
        x0, y0, x1, y1 = box_img
        crop = img.crop((int(math.floor(x0)), int(math.floor(y0)), int(math.ceil(x1)), int(math.ceil(y1))))
        cw, ch = crop.size
        if cw < 1 or ch < 1:
            return [], []
        U = max(1, int(round(48 * cw / ch)))
        strip = np.asarray(crop.resize((U, 48), OPTS["resample"]), dtype=np.float32)
        labels, probs = [], []
        starts = [0] if U <= 320 else list(range(0, U - 320 + 223, 224))
        starts = [s for s in starts if s == 0 or s < U]
        n_frames = -(-U // 4)
        for xs in starts:
            win = np.full((48, 320, 3), 127.0, dtype=np.float32)
            seg = strip[:, xs:xs + 320]
            win[:, : seg.shape[1]] = seg
            J = torch.from_numpy(win / 127.5 - 1.0).permute(2, 0, 1)[None]
            q = self.reader_frames(J)                            # [80, 307]
            lo = 0 if xs == 0 else 12
            hi = (-(-(U - xs) // 4)) if xs + 320 >= U else 68
            for t in range(lo, min(hi, 80)):
                if xs // 4 + t < n_frames:
                    labels.append(int(q[t].argmax()))
                    probs.append(q[t])
        return labels, probs

    def reader_frames(self, J):
        p = self.p
        G = cnorm(conv3(J, p["reader/w0"], (2, 2)), p["reader/stage_norm_0/scale"])
        for j, stride in enumerate(((2, 2), (2, 1), (2, 1))):
            y = dwconv(F.silu(G), p[f"reader/dw_{j}"], stride)
            G = cnorm(pw(y, p[f"reader/pw_{j}/kernel"]), p[f"reader/stage_norm_{j + 1}/scale"])
        Z = F.silu(G).mean(2)[0].transpose(0, 1)                   # [80, 384]
        for b in range(4):
            h = zc_norm(Z, p[f"reader/norm_{b}/scale"])
            k = p[f"reader/mix_{b}"][0, :, 0, :]                   # [5, 384]
            h = F.conv1d(F.pad(h.T[None], (2, 2)), k.T[:, None, :], groups=384)[0].T
            vg = h @ p[f"reader/up_{b}/kernel"]
            Z = Z + (vg[:, :384] * F.silu(vg[:, 384:])) @ p[f"reader/down_{b}/kernel"]
        logits = zc_norm(Z, p["reader/final_norm/scale"]) @ p["reader/out/kernel"] + p["reader/out/bias"]
        return logits.softmax(-1)

    # ------------------------------------------------------------ slots/memory
    def fourier(self, b):
        r = torch.arange(8, dtype=torch.float32)
        parts = []
        for c in range(4):
            a = (2.0 ** r) * math.pi * float(b[c])
            parts += [a.sin(), a.cos()]
        return torch.cat(parts)

    @torch.inference_mode()
    def memory(self, enc, objects, texts, vocab):
        p = self.p
        E, D, valid = enc["E"][0], enc["D"][0], enc["valid"]
        rows, cols = torch.arange(256) // 16, torch.arange(256) % 16
        cells = E + torch.sigmoid(p["pe_gate"]) * sincos_2d(rows, cols)
        slots, refs = [], []
        centres = (torch.arange(16) + 0.5) / 16
        for kind, boxes in (("object", objects), ("text", texts)):
            for k, bx in enumerate(boxes):
                b = bx["box"]
                i, j = bx["cell"]
                inside = ((centres[:, None] >= b[1]) & (centres[:, None] <= b[3])
                          & (centres[None, :] >= b[0]) & (centres[None, :] <= b[2])).flatten()
                ci = min(15, int((b[1] + b[3]) / 2 * 16)) * 16 + min(15, int((b[0] + b[2]) / 2 * 16))
                inside[ci] = True
                inside &= valid
                if not inside.any():
                    inside[ci] = True
                pooled = E[inside].mean(0)
                z = torch.cat([D[:, i, j], self.fourier(b)]) @ p["slot_proj/kernel"] + pooled @ p["slot_pool/kernel"]
                rid = vocab.ref_id(kind, k)
                slots.append(zc_norm(z, p["slot_norm/scale"]) + math.sqrt(D_MODEL) * self.T[rid])
                refs.append(rid)
        M = torch.cat([cells] + ([torch.stack(slots)] if slots else []), 0)
        mvalid = torch.cat([valid, torch.ones(len(slots), dtype=torch.bool)])
        return M, mvalid, refs


class Decoder:
    """Incremental decoder with self-attention, tap, engram and cross-attention caches."""

    def __init__(self, model: Mirage, M, mvalid):
        self.m, self.p = model, model.p
        p = self.p
        pre = "stack/layers/block/"
        self.ck, self.cv = [], []
        for l in range(8):
            k = (M @ p[pre + "cross_attn/k_proj/kernel"][l]).view(-1, 8, 48).transpose(0, 1)
            self.ck.append(zc_norm(k, p[pre + "cross_attn/k_norm/scale"][l]))
            self.cv.append((M @ p[pre + "cross_attn/v_proj/kernel"][l]).view(-1, 8, 64).transpose(0, 1))
        self.mvalid = mvalid
        self.tokens = []
        self.K = [[] for _ in range(8)]
        self.V = [[] for _ in range(8)]
        self.raw = [{"q": [], "k": [], "v": []} for _ in range(8)]
        self.engram_v = {s: [] for s in ENGRAM_SITES}

    def _engram(self, site_i, layer, t):
        p = self.p
        pre = f"engrams_{site_i}/"
        rows = []
        for oi, order in enumerate(ENGRAM_ORDERS):
            for h in range(ENGRAM_HEADS):
                tbl = p[pre + "embedding"][oi * ENGRAM_HEADS + h]
                if t - (order - 1) < 0:
                    rows.append(torch.zeros(tbl.shape[1]))
                else:
                    rows.append(tbl[engram_index(self.tokens, t, oi, order, h)])
        e = torch.cat(rows)
        k = e @ p[pre + "key_proj/kernel"]
        vraw = e @ p[pre + "value_proj/kernel"]
        hist = self.engram_v[layer]
        hist.append(vraw)
        taps = p[pre + "taps"]
        v = torch.zeros_like(vraw)
        for j in range(ENGRAM_TAPS):
            if t - ENGRAM_DILATION * j >= 0:
                v = v + taps[j] * hist[t - ENGRAM_DILATION * j]
        return k, v

    def _hmlp(self, x, l):
        p, pre = self.p, "stack/layers/block/hadamard_mlp/"

        def kron(z, a, b):
            return (a.T @ z.view(16, 32) @ b).reshape(-1)

        cond = 1 + torch.softmax(x @ p[pre + "cond_v"][l], -1) @ p[pre + "cond_u"][l]
        z = kron(p[pre + "d1"][l] * x, p[pre + "w1a"][l], p[pre + "w1b"][l])[self.m.perm1]
        z = kron(F.silu(p[pre + "d2"][l] * cond * z + p[pre + "b2"][l]), p[pre + "w2a"][l], p[pre + "w2b"][l])[self.m.perm2]
        z = kron(p[pre + "d3"][l] * z, p[pre + "w3a"][l], p[pre + "w3b"][l])
        return p[pre + "d4"][l] * z

    def _block(self, l, u, t):
        p, pre = self.p, "stack/layers/block/"
        x = u
        if l in ENGRAM_SITES:
            k, v = self._engram(ENGRAM_SITES.index(l), l, t)
            alpha = torch.sigmoid((rms_unit(x) * rms_unit(k)).sum() / math.sqrt(D_MODEL))
            x = x + alpha * v
        skip = x
        xn = zc_norm(x, p[pre + "ZCRMSNorm_0/scale"][l])
        R = self.raw[l]
        out = {}
        for name, n_heads, hd in (("q", 8, 48), ("k", 2, 48), ("v", 2, 64)):
            r = xn @ p[pre + f"self_attn/{name}_proj/kernel"][l]
            R[name].append(r)
            taps = p[pre + f"self_attn/{name}_taps"][l]
            acc = torch.zeros_like(r)
            for j in range(3):
                if t - j >= 0:
                    acc = acc + taps[j] * R[name][t - j]
            out[name] = acc.view(n_heads, hd)
        q = zc_norm(out["q"], p[pre + "self_attn/q_norm/scale"][l])
        k = zc_norm(out["k"], p[pre + "self_attn/k_norm/scale"][l])
        pos = torch.tensor(t)
        q, k = rope_half(q, pos, 100000.0), rope_half(k, pos, 100000.0)
        self.K[l].append(k)
        self.V[l].append(out["v"])
        Ks, Vs = torch.stack(self.K[l], 1), torch.stack(self.V[l], 1)     # [2, t+1, d]
        Ks, Vs = Ks.repeat_interleave(4, 0), Vs.repeat_interleave(4, 0)   # [8, t+1, d]
        att = torch.softmax((Ks @ q[:, :, None])[..., 0] / math.sqrt(48), -1)  # [8, t+1]
        o = (att[:, None, :] @ Vs)[:, 0, :].reshape(-1)
        A = (o * torch.sigmoid(xn @ p[pre + "self_attn/gate_proj/kernel"][l])) @ p[pre + "self_attn/out_proj/kernel"][l]
        x = skip + torch.sigmoid(p[pre + "attn_gate"][l]) * zc_norm(A, p[pre + "post_attn_norm/scale"][l])
        skip = x
        xc = zc_norm(x, p[pre + "cross_norm/scale"][l])
        qc = zc_norm((xc @ p[pre + "cross_attn/q_proj/kernel"][l]).view(8, 48), p[pre + "cross_attn/q_norm/scale"][l])
        s = (self.ck[l] @ qc[:, :, None])[..., 0] / math.sqrt(48)
        s = s.masked_fill(~self.mvalid[None], float("-inf"))
        oc = (torch.softmax(s, -1)[:, None, :] @ self.cv[l])[:, 0, :].reshape(-1)
        CA = (oc * torch.sigmoid(xc @ p[pre + "cross_attn/gate_proj/kernel"][l])) @ p[pre + "cross_attn/out_proj/kernel"][l]
        x = skip + torch.sigmoid(p[pre + "cross_gate"][l]) * zc_norm(CA, p[pre + "post_cross_norm/scale"][l])
        return x + self._hmlp(zc_norm(x, p[pre + "pre_hada_norm/scale"][l]), l)

    @torch.inference_mode()
    def step(self, tok):
        p = self.p
        t = len(self.tokens)
        self.tokens.append(int(tok))
        X = (math.sqrt(D_MODEL) * self.m.T[tok]).expand(4, D_MODEL).clone()
        for l in range(8):
            nx = rms_unit(X.reshape(-1))
            lane = torch.zeros(4)
            lane[l % 4] = 1.0
            hpre = torch.sigmoid(p["stack/mhc_a_pre"][l] * (nx @ p["stack/mhc_phi_pre"][l]) + p["stack/mhc_b_pre"][l] + 8 * lane - 4)
            u = (hpre[:, None] * X).sum(0)
            y = self._block(l, u, t) - u
            hpost = 2 * torch.sigmoid(p["stack/mhc_a_post"][l] * (nx @ p["stack/mhc_phi_post"][l]) + p["stack/mhc_b_post"][l] - 4 * (1 - lane))
            hres = sinkhorn(p["stack/mhc_a_res"][l] * (nx @ p["stack/mhc_phi_res"][l]).view(4, 4) + p["stack/mhc_b_res"][l])
            X = hres @ X + hpost[:, None] * y
        h = zc_norm(X.mean(0), p["stack/final_norm/scale"])
        return h @ self.m.T.T


# Full-sequence (teacher-forced) decoder: logits for every position of a token sequence in one pass,
# equivalent to Decoder stepped token by token; for training and for scoring weight variants.
PRE = "stack/layers/block/"


def shift(x, j):
    """Row t of the result is row t - j of x (zeros before the start)."""
    if j == 0:
        return x
    return torch.cat([x.new_zeros((j,) + x.shape[1:]), x[:-j]], 0) if j < x.shape[0] else torch.zeros_like(x)


def engram_rows(p, tokens, site_i):
    T = len(tokens)
    tok = [int(t) for t in tokens]
    parts = []
    for oi, order in enumerate(ENGRAM_ORDERS):
        for h in range(ENGRAM_HEADS):
            tbl = p[f"engrams_{site_i}/embedding"][oi * ENGRAM_HEADS + h]
            idx = torch.tensor([engram_index(tok, t, oi, order, h) for t in range(T)])
            rows = tbl[idx]
            ok = (torch.arange(T) >= order - 1).float()[:, None]
            parts.append(rows * ok)
    return torch.cat(parts, -1)


def hmlp(p, x, l, perm1, perm2):
    pre = PRE + "hadamard_mlp/"

    def kron(z, a, b):
        return torch.einsum("tij,ik,jl->tkl", z.reshape(z.shape[0], 16, 32), a, b).reshape(z.shape[0], -1)

    cond = 1 + torch.softmax(x @ p[pre + "cond_v"][l], -1) @ p[pre + "cond_u"][l]
    z = kron(p[pre + "d1"][l] * x, p[pre + "w1a"][l], p[pre + "w1b"][l])[:, perm1]
    z = kron(F.silu(p[pre + "d2"][l] * cond * z + p[pre + "b2"][l]), p[pre + "w2a"][l], p[pre + "w2b"][l])[:, perm2]
    z = kron(p[pre + "d3"][l] * z, p[pre + "w3a"][l], p[pre + "w3b"][l])
    return p[pre + "d4"][l] * z


def cross_kv(p, M):
    ck, cv = [], []
    for l in range(8):
        k = (M @ p[PRE + "cross_attn/k_proj/kernel"][l]).view(-1, 8, 48).transpose(0, 1)
        ck.append(zc_norm(k, p[PRE + "cross_attn/k_norm/scale"][l]))
        cv.append((M @ p[PRE + "cross_attn/v_proj/kernel"][l]).view(-1, 8, 64).transpose(0, 1))
    return ck, cv


def decoder_forward(model, M, mvalid, tokens):
    """tokens: 1-D LongTensor. Returns logits [T, vocab]; row t predicts token t + 1."""
    p = model.p
    T = tokens.shape[0]
    ck, cv = cross_kv(p, M)
    X = (math.sqrt(D_MODEL) * model.T[tokens])[:, None, :].expand(T, 4, D_MODEL)
    pos = torch.arange(T)
    causal = torch.ones(T, T, dtype=torch.bool).tril()
    eng = {s: engram_rows(p, tokens, i) for i, s in enumerate(ENGRAM_SITES)}
    for l in range(8):
        nx = rms_unit(X.reshape(T, -1))
        lane = torch.zeros(4)
        lane[l % 4] = 1.0
        hpre = torch.sigmoid(p["stack/mhc_a_pre"][l] * (nx @ p["stack/mhc_phi_pre"][l]) + p["stack/mhc_b_pre"][l] + 8 * lane - 4)
        u = (hpre[..., None] * X).sum(1)
        x = u
        if l in ENGRAM_SITES:
            i = ENGRAM_SITES.index(l)
            e = eng[l]
            k = e @ p[f"engrams_{i}/key_proj/kernel"]
            vraw = e @ p[f"engrams_{i}/value_proj/kernel"]
            taps = p[f"engrams_{i}/taps"]
            v = sum(taps[j] * shift(vraw, ENGRAM_DILATION * j) for j in range(ENGRAM_TAPS))
            alpha = torch.sigmoid((rms_unit(x) * rms_unit(k)).sum(-1, keepdim=True) / math.sqrt(D_MODEL))
            x = x + alpha * v
        skip = x
        xn = zc_norm(x, p[PRE + "ZCRMSNorm_0/scale"][l])
        qkv = {}
        for name, heads, hd in (("q", 8, 48), ("k", 2, 48), ("v", 2, 64)):
            raw = xn @ p[PRE + f"self_attn/{name}_proj/kernel"][l]
            taps = p[PRE + f"self_attn/{name}_taps"][l]
            qkv[name] = sum(taps[j] * shift(raw, j) for j in range(3)).view(T, heads, hd)
        q = rope_half(zc_norm(qkv["q"], p[PRE + "self_attn/q_norm/scale"][l]), pos[:, None], 100000.0)
        k = rope_half(zc_norm(qkv["k"], p[PRE + "self_attn/k_norm/scale"][l]), pos[:, None], 100000.0)
        q, k, v = q.transpose(0, 1), k.transpose(0, 1).repeat_interleave(4, 0), qkv["v"].transpose(0, 1).repeat_interleave(4, 0)
        s = (q @ k.transpose(-1, -2)) / math.sqrt(48)
        o = (s.masked_fill(~causal, float("-inf")).softmax(-1) @ v).transpose(0, 1).reshape(T, -1)
        A = (o * torch.sigmoid(xn @ p[PRE + "self_attn/gate_proj/kernel"][l])) @ p[PRE + "self_attn/out_proj/kernel"][l]
        x = skip + torch.sigmoid(p[PRE + "attn_gate"][l]) * zc_norm(A, p[PRE + "post_attn_norm/scale"][l])
        skip = x
        xc = zc_norm(x, p[PRE + "cross_norm/scale"][l])
        qc = zc_norm((xc @ p[PRE + "cross_attn/q_proj/kernel"][l]).view(T, 8, 48), p[PRE + "cross_attn/q_norm/scale"][l]).transpose(0, 1)
        sc = (qc @ ck[l].transpose(-1, -2)) / math.sqrt(48)
        oc = (sc.masked_fill(~mvalid[None, None], float("-inf")).softmax(-1) @ cv[l]).transpose(0, 1).reshape(T, -1)
        CA = (oc * torch.sigmoid(xc @ p[PRE + "cross_attn/gate_proj/kernel"][l])) @ p[PRE + "cross_attn/out_proj/kernel"][l]
        x = skip + torch.sigmoid(p[PRE + "cross_gate"][l]) * zc_norm(CA, p[PRE + "post_cross_norm/scale"][l])
        x = x + hmlp(p, zc_norm(x, p[PRE + "pre_hada_norm/scale"][l]), l, model.perm1, model.perm2)
        y = x - u
        hpost = 2 * torch.sigmoid(p["stack/mhc_a_post"][l] * (nx @ p["stack/mhc_phi_post"][l]) + p["stack/mhc_b_post"][l] - 4 * (1 - lane))
        hres = sinkhorn(p["stack/mhc_a_res"][l] * (nx @ p["stack/mhc_phi_res"][l]).view(T, 4, 4) + p["stack/mhc_b_res"][l])
        X = hres @ X + hpost[..., None] * y[:, None, :]
    h = zc_norm(X.mean(1), p["stack/final_norm/scale"])
    return h @ model.T.T


# ===================================================================================== text reader
# Trainable, batched copy of Mirage's convolutional CTC text reader. Input: [B, 3, 48, 320] strips scaled to
# [-1, 1]. Output: [B, 80, 307] logits, class 0 the CTC blank, class c > 0 the charset entry c - 1.
PREFIX = "reader/"
QUANTISED = {"pw_0/kernel", "pw_1/kernel", "pw_2/kernel", "out/kernel"} | \
    {f"{k}_{b}/kernel" for k in ("up", "down") for b in range(4)}


class Reader(nn.Module):
    """quantizer: a Quantizer; when given, the 2-bit matrices run through fake quantisation."""

    def __init__(self, weights, quantizer=None, bits=2, act_quant=False):
        super().__init__()
        self.names = sorted(k for k in weights if k.startswith(PREFIX))
        self.params = nn.ParameterDict({k.replace("/", "__"): nn.Parameter(weights[k].clone().float()) for k in self.names})
        self.quantizer, self.bits, self.act_quant = quantizer, bits, act_quant

    def aq(self, x):
        """Engine-style activation quantisation: each 128-group rotated, int8 per group, rotated back."""
        if not self.act_quant or self.quantizer is None:
            return x
        d = x.shape[-1]
        pc = -(-d // 128) * 128
        with torch.no_grad():
            xp = F.pad(x, (0, pc - d)) if pc != d else x
            g = xp.reshape(*xp.shape[:-1], pc // 128, 128) @ self.quantizer.H
            s = g.abs().amax(-1, keepdim=True).clamp_min(1e-12) / 127.0
            q = (torch.round(g / s) * s) @ self.quantizer.H
            q = q.reshape(xp.shape)[..., :d]
        return x + (q - x).detach()

    def w(self, name):
        p = self.params[(PREFIX + name).replace("/", "__")]
        if self.quantizer is not None and name in QUANTISED:
            return self.quantizer.fake_quant_kernel(p, self.bits)
        return p

    def forward(self, J):
        G = cnorm(conv3(J, self.w("w0"), (2, 2)), self.w("stage_norm_0/scale"))
        for j, stride in enumerate(((2, 2), (2, 1), (2, 1))):
            y = dwconv(F.silu(G), self.w(f"dw_{j}"), stride)
            G = cnorm((self.aq(y.permute(0, 2, 3, 1)) @ self.w(f"pw_{j}/kernel")).permute(0, 3, 1, 2), self.w(f"stage_norm_{j + 1}/scale"))
        Z = F.silu(G).mean(2).transpose(1, 2)                         # [B, 80, 384]
        for b in range(4):
            h = zc_norm(Z, self.w(f"norm_{b}/scale"))
            k = self.w(f"mix_{b}")[0, :, 0, :]                        # [5, 384]
            h = F.conv1d(F.pad(h.transpose(1, 2), (2, 2)), k.T[:, None, :], groups=384).transpose(1, 2)
            vg = self.aq(h) @ self.w(f"up_{b}/kernel")
            Z = Z + self.aq(vg[..., :384] * F.silu(vg[..., 384:])) @ self.w(f"down_{b}/kernel")
        return self.aq(zc_norm(Z, self.w("final_norm/scale"))) @ self.w("out/kernel") + self.w("out/bias")

    def state(self):
        return {k: self.params[k.replace("/", "__")].detach().cpu().contiguous() for k in self.names}


def greedy(logits, charset, lengths=None):
    ids = logits.argmax(-1).cpu()
    out = []
    for b in range(ids.shape[0]):
        n = ids.shape[1] if lengths is None else int(lengths[b])
        out.append(ctc_text(ids[b, :n].tolist(), charset))
    return out


# ===================================================================================== record assembly
# Turns the decoder's compact lines plus boxes, masks and OCR strings into one JSON record. Decoder grammar,
# as emitted by the model:
#   <think>trace</think>
#   S key=value; key=value
#   E<n> name @position *<objK><objJ> {key=value; key=value}
#   T *<txtK>
#   R <subject n> predicate <object n>
# Unrecognised lines are kept under "unparsed".
REF = re.compile(r"<(obj|txt)(\d+)>")
ENTITY = re.compile(r"^E(\d+)\s+(.*)$")


def _attrs(s):
    out = {}
    for part in s.split(";"):
        if "=" in part:
            k, v = part.split("=", 1)
            out[k.strip()] = v.strip()
        elif part.strip():
            out.setdefault("note", part.strip())
    return out


def parse_decoded(text):
    trace = ""
    body = text
    m = re.search(r"<think>(.*?)</think>", text, flags=re.S)
    if m:
        trace = m.group(1).strip()
        body = text[m.end():]
    body = body.replace("</s>", "")
    rec = dict(scene={}, entities=[], text_refs=[], relations=[], events=[], unparsed=[], reasoning=trace)
    for line in body.split("\n"):
        line = line.strip()
        if not line:
            continue
        if line.startswith("S "):
            rec["scene"].update(_attrs(line[2:]))
            continue
        em = ENTITY.match(line)
        if em:
            n, rest = int(em.group(1)), em.group(2)
            attrs = {}
            am = re.search(r"\{(.*)\}\s*$", rest)
            if am:
                attrs = _attrs(am.group(1))
                rest = rest[:am.start()].strip()
            refs = [(k, int(i)) for k, i in REF.findall(rest)]
            rest = REF.sub("", rest).replace("*", " ").strip()
            pos = None
            if "@" in rest:
                rest, pos = rest.split("@", 1)
                pos = pos.strip()
            name, count = rest.strip(), None
            cm = re.search(r"\s*#(\S+)$", name)
            if cm:
                name, count = name[:cm.start()].strip(), cm.group(1)
            ent = dict(n=n, name=name, position=pos, refs=refs, description=attrs)
            if count is not None:
                ent["count"] = int(count) if count.isdigit() else count
            rec["entities"].append(ent)
            continue
        if line.startswith("T"):
            rec["text_refs"] += [int(i) for k, i in REF.findall(line) if k == "txt"]
            continue
        rm = re.match(r"^R\s+(\d+)\s+(\S+)\s+(\d+)\s*$", line)
        if rm:
            rec["relations"].append(dict(subject=int(rm.group(1)), predicate=rm.group(2), object=int(rm.group(3))))
            continue
        if line.startswith("V ") or line.startswith("EV"):
            rec["events"].append(line.split(" ", 1)[1] if " " in line else line)
            continue
        rec["unparsed"].append(line)
    return rec


def assemble(decoded, objects, texts, ocr, to_1000, masks=None, fields=None):
    """objects/texts: detector proposals (score, box); ocr: strings per text proposal; to_1000: box mapper."""
    rec = parse_decoded(decoded)
    out = dict(scene=rec["scene"], objects=[], regions=[], text=[], relations=rec["relations"],
               events=rec["events"], reasoning=rec["reasoning"])
    for e in rec["entities"]:
        item = dict(id=e["n"], name=e["name"])
        if "count" in e:
            item["count"] = e["count"]
        if e["position"]:
            item["position"] = e["position"]
        if e["description"]:
            item["description"] = e["description"]
        objs = [i for k, i in e["refs"] if k == "obj" and i < len(objects)]
        if objs:
            item["bbox"] = to_1000(objects[objs[0]]["box"])
            item["score"] = round(objects[objs[0]]["score"], 3)
            if len(objs) > 1:
                item["count"] = len(objs)
                item["instances"] = [dict(bbox=to_1000(objects[i]["box"]), score=round(objects[i]["score"], 3)) for i in objs]
            if masks is not None and objs[0] in masks:
                item["mask"] = masks[objs[0]]
            out["objects"].append(item)
        else:
            out["regions"].append(item)
    for k, (b, s) in enumerate(zip(texts, ocr)):
        if not s:
            continue
        out["text"].append(dict(string=s, bbox=to_1000(b["box"]), score=round(b["score"], 3),
                                mentioned=k in rec["text_refs"]))
    if rec["unparsed"]:
        out["unparsed"] = rec["unparsed"]
    return out


class Pipeline:
    """describe(): image in, structured record out, using the reference model."""

    def __init__(self, weights, archive, reader_weights=None):
        self.model = Mirage(weights)
        if reader_weights:
            for k, v in load_file(reader_weights).items():
                self.model.p[k] = v.float()
        arc = Archive(archive)
        self.vocab = vocab_of(arc, self.model.T.shape[0])
        self.charset = charset_of(arc)

    def _ocr(self, img, box_1000, meta):
        px = (box_1000[0] / 1000 * meta["W"], box_1000[1] / 1000 * meta["H"],
              box_1000[2] / 1000 * meta["W"], box_1000[3] / 1000 * meta["H"])
        labels, probs = self.model.read(img, px)
        return ctc_text(labels, self.charset)

    def _polygons(self, prob, meta, fmt):
        import cv2
        cv2.setNumThreads(1)
        m = (prob.numpy() > 0.5).astype(np.uint8)
        if not m.any():
            return []
        contours, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        scale = CANVAS / 128 / meta["r"]
        rings = []
        for c in sorted(contours, key=cv2.contourArea, reverse=True):
            c = cv2.approxPolyDP(c, 0.6, True)[:, 0, :].astype(np.float64)
            if len(c) < 3:
                continue
            xs = np.clip((c[:, 0] + 0.5) * scale / meta["W"] * 1000, 0, 1000)
            ys = np.clip((c[:, 1] + 0.5) * scale / meta["H"] * 1000, 0, 1000)
            rings.append([int(round(v)) for xy in zip(xs, ys) for v in xy])
        return rings[:1][0] if fmt == "largest" and rings else rings

    @torch.inference_mode()
    def describe(self, image, language="en", conf=0.5, text_conf=0.35, max_objects=20, max_text_lines=20,
                 max_tokens=512, masks=True, mask_format="largest"):
        t0 = time.time()
        m = self.model
        enc = m.encode(image)
        meta = enc["meta"]
        det = m.detect(enc, conf=conf, text_conf=text_conf, nms_iou=0.5)
        objects, texts = det["object"][:max_objects], det["text"][:max_text_lines]
        to_1000 = lambda b: m.to_1000(b, meta)  # noqa: E731
        ocr = [self._ocr(image, to_1000(b["box"]), meta) for b in texts]
        t1 = time.time()
        M, mvalid, refs = m.memory(enc, objects, texts, self.vocab)
        dec = Decoder(m, M, mvalid)
        v = self.vocab
        banned = torch.zeros(v.size, dtype=torch.bool)
        banned[list(v.lang_ids.values())] = True
        banned[v.ref_base:] = True
        banned[refs] = False
        logits = None
        for tid in (v.bos, v.lang_ids[language]):
            logits = dec.step(tid)
        gen = []
        for _ in range(max_tokens):
            nxt = int(logits.masked_fill(banned, float("-inf")).argmax())
            if nxt == v.eos:
                break
            gen.append(nxt)
            logits = dec.step(nxt)
        t2 = time.time()
        decoded = v.decode(gen)
        mask_out = None
        if masks and objects:
            probs = m.masks(enc, objects)
            mask_out = {i: self._polygons(p, meta, mask_format) for i, p in enumerate(probs)}
        out = assemble(decoded, objects, texts, ocr, to_1000, masks=mask_out)
        out["timing_ms"] = dict(perception=round(1000 * (t1 - t0)), decode=round(1000 * (t2 - t1)),
                                tokens=len(gen), total=round(1000 * (time.time() - t0)))
        return out

# ===================================================================================== synthetic text
# Synthetic text-line strips for the CTC reader: sampled strings rendered with system fonts and augmented.
# Each sample is a [48, 320, 3] uint8 strip (text left-aligned, grey 127 padding), the used width U and the
# label ids (charset index + 1).
STRIP_H, STRIP_W = 48, 320

# Accented letters and marks of the seven description languages (en, de, fr, es, it, nl, pl).
COMMON = set("àâäæçéèêëîïôœùûüÿÀÂÄÆÇÉÈÊËÎÏÔŒÙÛÜŸßáíñóúÁÍÑÓÚìòÌÒijIJąćęłńśźżĄĆĘŁŃŚŹŻ¡¿«»–—")
# Letters found only in romanisations and obsolete orthographies.
RARE = set("ĔĕĬĭŎŏĸʼnĿŀ")


def letter_weight(c):
    return 4.0 if c in COMMON else 0.2 if c in RARE else 1.0


def font_paths(patterns):
    paths = []
    for p in patterns:
        paths += glob.glob(p)
    return sorted(set(paths))


class FontPool:
    def __init__(self, paths, charset):
        self.paths, self.cover = [], []
        probe = "\U0010fffd"
        for path in paths:
            try:
                f = ImageFont.truetype(path, 32)
            except OSError:
                continue
            missing = np.asarray(f.getmask(probe))
            ok = set()
            for c in charset:
                if c == " ":
                    ok.add(c)
                    continue
                m = f.getmask(c)
                if m.size[0] == 0:
                    continue
                arr = np.asarray(m)
                if arr.shape == missing.shape and np.array_equal(arr, missing):
                    continue
                ok.add(c)
            self.paths.append(path)
            self.cover.append(ok)
        self.cache = {}

    def get(self, rng, text, size):
        need = set(text)
        cands = [i for i, c in enumerate(self.cover) if need <= c]
        if not cands:
            return None
        i = rng.choice(cands)
        key = (i, size)
        if key not in self.cache:
            self.cache[key] = ImageFont.truetype(self.paths[i], size)
        return self.cache[key]


class TextSampler:
    def __init__(self, charset, dead, words, seed=0):
        self.charset = charset
        self.chars = [c for c in charset if c != " "]
        self.dead = [c for c in dead if c in charset]
        self.words = [w for w in words if w and all(ch in charset for ch in w)]
        self.variants = {}
        for c in self.chars:
            base = unicodedata.normalize("NFD", c)[0]
            if base != c and base.isalpha():
                self.variants.setdefault(base, []).append(c)
        self.dead_by_base = {}
        for c in self.dead:
            base = unicodedata.normalize("NFD", c)[0]
            self.dead_by_base.setdefault(base, []).append(c)
        self.letters_with_dead = [b for b in self.dead_by_base]
        self.rng = random.Random(seed)

    def _accent(self, word, p_dead):
        r = self.rng
        out = []
        for ch in word:
            if ch in self.dead_by_base and r.random() < p_dead:
                out.append(r.choice(self.dead_by_base[ch]))
            elif ch in self.variants and r.random() < 0.15:
                out.append(r.choice(self.variants[ch]))
            else:
                out.append(ch)
        return "".join(out)

    def _case(self, s):
        r = self.rng.random()
        if r < 0.3:
            return s.upper()
        if r < 0.45:
            return s.title()
        if r < 0.6:
            return s.lower()
        return s

    def sample(self):
        r = self.rng
        mode = r.random()
        if mode < 0.40:
            s = " ".join(self._case(self._accent(r.choice(self.words), 0.08)) for _ in range(r.randint(1, 3)))
        elif mode < 0.60:
            pool = self.chars + self.dead * 4
            s = "".join(r.choice(pool) for _ in range(r.randint(2, 12)))
        elif mode < 0.75:
            words = []
            for _ in range(r.randint(1, 3)):
                w = list(self._case(r.choice(self.words)))
                if not w:
                    continue
                k = r.randrange(len(w))
                w[k] = r.choice(self.dead)
                words.append("".join(w))
            s = " ".join(words) or r.choice(self.dead)
        elif mode < 0.85:
            num = f"{r.randint(0, 9999)}{r.choice(['', ',', '.'])}{r.randint(0, 99):02d}"
            s = r.choice([
                f"{num} {r.choice('€£¥$')}", f"{r.choice('€£¥$')}{num}", f"{r.randint(1, 31)}–{r.randint(1, 31)}",
                f"¡{self._case(r.choice(self.words))}!", f"¿{self._case(r.choice(self.words))}?",
                f"«{r.choice(self.words)}»", f"“{r.choice(self.words)}”", f"{r.choice(self.words)} – {r.choice(self.words)}",
                f"{r.randint(1, 12)}:{r.randint(0, 59):02d}", f"{r.randint(0, 100)}%", f"№{num}" if "№" in self.charset else f"#{num}",
            ])
        else:
            s = " ".join(self._accent(r.choice(self.words), 0.15).upper() for _ in range(r.randint(1, 3)))
        s = "".join(ch for ch in s if ch in self.charset).strip()
        return s or r.choice(self.dead)


class RealisticSampler(TextSampler):
    """Text with letter frequencies close to real writing in the charset's languages.

    Real words keep their own accents; contrast lines swap one letter for a sibling with the same base
    letter (Ê for E, Ĕ for E, ...) weighted by how common it is, so accents have to be read rather than
    guessed; a small share of random strings covers every class.
    """

    def __init__(self, charset, dead, words, seed=0):
        super().__init__(charset, dead, words, seed)
        self.siblings = {}
        for c in self.chars:
            base = unicodedata.normalize("NFD", c)[0]
            if base.isalpha() and base in charset:
                self.siblings.setdefault(base, []).append(c)
        self.siblings = {b: [c for c in v if c != b] for b, v in self.siblings.items()}
        self.siblings = {b: v for b, v in self.siblings.items() if v}
        self.pool = self.chars
        self.pool_w = [letter_weight(c) for c in self.chars]

    def _swap(self, word):
        spots = [i for i, ch in enumerate(word) if ch in self.siblings or unicodedata.normalize("NFD", ch)[0] in self.siblings]
        if not spots:
            return word
        i = self.rng.choice(spots)
        base = unicodedata.normalize("NFD", word[i])[0]
        opts = self.siblings[base]
        c = self.rng.choices(opts, weights=[letter_weight(o) for o in opts])[0]
        return word[:i] + c + word[i + 1:]

    def sample(self):
        r = self.rng
        mode = r.random()
        if mode < 0.45:
            s = " ".join(self._case(r.choice(self.words)) for _ in range(r.randint(1, 3)))
        elif mode < 0.80:
            words = [self._case(r.choice(self.words)) for _ in range(r.randint(1, 3))]
            k = r.randrange(len(words))
            words[k] = self._swap(words[k])
            s = " ".join(words)
        elif mode < 0.90:
            s = "".join(r.choices(self.pool, weights=self.pool_w, k=r.randint(2, 12)))
        else:
            return super().sample()
        s = "".join(ch for ch in s if ch in self.charset).strip()
        return s or r.choice(self.chars)


def _contrast_colors(rng):
    while True:
        a = np.array([rng.randint(0, 255) for _ in range(3)])
        b = np.array([rng.randint(0, 255) for _ in range(3)])
        la, lb = a @ [0.299, 0.587, 0.114], b @ [0.299, 0.587, 0.114]
        if abs(la - lb) > 70:
            return tuple(int(x) for x in a), tuple(int(x) for x in b)


def _background(rng, w, h, color):
    kind = rng.random()
    base = np.zeros((h, w, 3), np.float32) + np.array(color, np.float32)
    if kind < 0.4:
        pass
    elif kind < 0.65:
        other = np.array([rng.randint(0, 255) for _ in range(3)], np.float32)
        t = np.linspace(0, 1, w if rng.random() < 0.5 else h)
        t = t[None, :, None] if t.shape[0] == w else t[:, None, None]
        base = base * (1 - 0.35 * t) + other * 0.35 * t
    elif kind < 0.85:
        base = base + np.random.default_rng(rng.randint(0, 2**31)).normal(0, rng.uniform(5, 25), (h, w, 1))
    else:
        period = rng.randint(4, 20)
        stripes = ((np.arange(w) // period) % 2)[None, :, None] * rng.uniform(10, 30)
        base = base + stripes
    return Image.fromarray(np.clip(base, 0, 255).astype(np.uint8))


def render(text, fonts, rng, tight=0.0):
    """tight: share of lines cut as closely as a text detector's box, so accents can touch the edge."""
    for _ in range(4):
        size = rng.randint(20, 64)
        font = fonts.get(rng, text, size)
        if font is None:
            return None
        x0, y0, x1, y1 = font.getbbox(text)
        tw, th = x1 - x0, y1 - y0
        if tw <= 0 or th <= 0:
            return None
        if rng.random() < tight:
            mx, my = int(size * rng.uniform(0.0, 0.12)), int(size * rng.uniform(-0.03, 0.04))
        else:
            mx, my = int(size * rng.uniform(0.05, 0.35)), int(size * rng.uniform(0.05, 0.3))
        W, H = tw + 2 * mx, th + 2 * my
        U = max(1, round(STRIP_H * W / H))
        if U <= STRIP_W:
            break
        text = text[: max(1, len(text) * STRIP_W // U - 1)].strip() or text[:1]
    else:
        return None
    fg, bg = _contrast_colors(rng)
    img = _background(rng, W, H, bg)
    ImageDraw.Draw(img).text((mx - x0, my - y0), text, font=font, fill=fg)
    if rng.random() < 0.3:
        img = img.rotate(rng.uniform(-2.5, 2.5), resample=Image.BILINEAR, expand=False, fillcolor=bg)
    if rng.random() < 0.3:
        img = img.filter(ImageFilter.GaussianBlur(rng.uniform(0.3, 1.2)))
    if rng.random() < 0.2:
        f = rng.uniform(0.35, 0.7)
        small = img.resize((max(1, int(W * f)), max(1, int(H * f))), Image.BILINEAR)
        img = small.resize((W, H), Image.BILINEAR)
    if rng.random() < 0.3:
        buf = io.BytesIO()
        img.save(buf, "JPEG", quality=rng.randint(30, 90))
        img = Image.open(io.BytesIO(buf.getvalue())).convert("RGB")
    return text, img


def to_strip(img):
    W, H = img.size
    U = max(1, round(STRIP_H * W / H))
    if U > STRIP_W:
        return None, 0
    arr = np.full((STRIP_H, STRIP_W, 3), 127, np.uint8)
    arr[:, :U] = np.asarray(img.resize((U, STRIP_H), Image.BILINEAR))
    return arr, U


class Generator:
    def __init__(self, charset, dead, words, font_globs, seed, style="uniform", tight=0.0):
        self.charset = charset
        self.index = {c: i + 1 for i, c in enumerate(charset)}
        self.fonts = FontPool(font_paths(font_globs), charset)
        self.text = (RealisticSampler if style == "realistic" else TextSampler)(charset, dead, words, seed)
        self.rng = random.Random(seed + 1)
        self.tight = tight

    def sample(self):
        while True:
            s = self.text.sample()
            out = render(s, self.fonts, self.rng, self.tight)
            if out is None:
                continue
            s, img = out
            strip, U = to_strip(img)
            if strip is None or not s:
                continue
            labels = [self.index[c] for c in s]
            frames = -(-U // 4)
            reps = sum(1 for a, b in zip(labels, labels[1:]) if a == b)
            if frames < len(labels) + reps:
                continue
            return strip, U, labels, s


# ===================================================================================== reader evaluation
# Word crops packed as archives of strips 48 pixels high (pix [48, total, 3], offsets, labels, sources),
# read with the engine's window and frame-retention rules.
def reader_windows(strip):
    U = strip.shape[1]
    starts = [0] if U <= 320 else list(range(0, U - 320 + 223, 224))
    wins, keep = [], []
    n_frames = -(-U // 4)
    for xs in starts:
        win = np.full((48, 320, 3), 127, np.uint8)
        seg = strip[:, xs:xs + 320]
        win[:, :seg.shape[1]] = seg
        lo = 0 if xs == 0 else 12
        hi = -(-(U - xs) // 4) if xs + 320 >= U else 68
        keep.append([t for t in range(lo, min(hi, 80)) if xs // 4 + t < n_frames])
        wins.append(win)
    return wins, keep


def load_packed(path):
    """The windows of a packed crop archive, with each window's word, its retained frames, the labels and sources."""
    d = np.load(path, allow_pickle=False)
    pix, offs, labels, sources = d["pix"], d["offsets"], d["labels"], d["sources"]
    all_w, owner, keeps = [], [], []
    for i in range(len(labels)):
        ws, ks = reader_windows(pix[:, offs[i]:offs[i + 1]])
        all_w += ws
        owner += [i] * len(ws)
        keeps += ks
    return torch.from_numpy(np.stack(all_w)).permute(0, 3, 1, 2), owner, keeps, labels, sources


@torch.no_grad()
def reader_paths(net, X, owner, keeps, n, dev):
    """The retained frame labels of every word, reading the windows X in batches of 512."""
    ids = []
    for b in range(0, X.shape[0], 512):
        ids.append(net(X[b:b + 512].to(dev).float() / 127.5 - 1.0).argmax(-1).cpu())
    ids = torch.cat(ids).numpy() if ids else np.zeros((0, 80), np.int64)
    paths = [[] for _ in range(n)]
    for k, (o, kp) in enumerate(zip(owner, keeps)):
        paths[o] += ids[k, kp].tolist()
    return paths


@torch.no_grad()
def evaluate_reader(net, sets, charset, dead_chars, dev):
    was = net.training
    net.eval()
    out = {}
    for X, owner, keeps, labels, sources in sets:
        for i, path in enumerate(reader_paths(net, X, owner, keeps, len(labels), dev)):
            pred, gt = ctc_text(path, charset), str(labels[i])
            st = out.setdefault(str(sources[i]), dict(n=0, acc=0, exact=0, ed=0, len=0, dh=0, dt=0, fp=0, ah=0, at=0))
            st["n"] += 1
            st["acc"] += word_key(pred) == word_key(gt)
            st["exact"] += pred == gt
            st["ed"] += edit_distance(pred, gt)
            st["len"] += len(gt)
            for ch in set(gt) & dead_chars:
                st["dt"] += gt.count(ch)
                st["dh"] += min(gt.count(ch), pred.count(ch))
            for ch in set(gt):
                if ord(ch) > 127 and ch.isalpha():
                    st["at"] += gt.count(ch)
                    st["ah"] += min(gt.count(ch), pred.count(ch))
            st["fp"] += sum(1 for ch in pred if ch in dead_chars and ch not in gt)
    if was:
        net.train()
    return {src: dict(n=st["n"], word_acc=round(st["acc"] / st["n"], 4), exact=round(st["exact"] / st["n"], 4),
                      cer=round(st["ed"] / max(st["len"], 1), 4), dead_recall=round(st["dh"] / max(st["dt"], 1), 4),
                      accent_recall=round(st["ah"] / max(st["at"], 1), 4), dead_false=st["fp"]) for src, st in out.items()}


# ===================================================================================== commands: model
def pack_indices(idx, bits):
    flat = idx.reshape(-1).astype(np.uint8)
    per = 8 // bits
    flat = np.concatenate([flat, np.zeros((-len(flat)) % per, np.uint8)])
    out = np.zeros(len(flat) // per, np.uint8)
    for k in range(per):
        out |= (flat[k::per] << (bits * k)).astype(np.uint8)
    return out


def cq_blob_from_latent(R, bits, quantizer):
    """R: [out, in] float32. Same arithmetic as Quantizer (the training forward)."""
    rows, cols = R.shape
    pc = -(-cols // 128) * 128
    Rp = torch.nn.functional.pad(R.to(quantizer.H.device), (0, pc - cols))
    g = Rp.reshape(rows, pc // 128, 128) @ quantizer.H
    n = g.norm(dim=-1, keepdim=True)
    idx = torch.bucketize((g / n.clamp_min(1e-30)).contiguous(), quantizer.mid[bits])
    return pack_indices(idx.cpu().numpy(), bits).tobytes() + n.half().cpu().numpy().reshape(rows, -1).tobytes()


def compact_scales(blob, rows, cols, bits):
    """Re-encode a cq blob's fp16 group scales as 8-bit log codes with the two f32 bounds appended."""
    pc = -(-cols // 128) * 128
    ib = rows * pc * bits // 8
    nrm = np.frombuffer(blob, np.float16, count=rows * (pc // 128), offset=ib).astype(np.float64)
    nz = nrm[nrm > 0]
    lo, hi = (np.float32(np.log(nz.min())), np.float32(np.log(nz.max()))) if len(nz) else (np.float32(0), np.float32(0))
    span = float(hi) - float(lo)
    pos = np.log(np.where(nrm > 0, nrm, 1.0)) - float(lo)
    steps = np.round(pos / span * 254.0) if span > 0 else np.zeros_like(pos)
    code = np.where(nrm > 0, np.clip(steps, 0, 254) + 1, 0).astype(np.uint8)
    return bytes(blob[:ib]) + code.tobytes() + np.array([lo, hi], np.float32).tobytes()


def int_rows(a, bits, keep=None):
    """[rows, cols] float -> signed integers of `bits` bits in one LSB-first bit stream, then one f32 scale
    per row (the largest magnitude of a row maps to the largest integer). Entries outside the boolean mask
    `keep` are stored as zero."""
    a = np.asarray(a, np.float64).reshape(a.shape[0], -1)
    top = 2 ** (bits - 1) - 1
    s = np.abs(a).max(1) / top
    q = np.round(a / np.where(s > 0, s, 1.0)[:, None]).clip(-top, top).astype(np.int64)
    if keep is not None:
        q = np.where(keep, q, 0)
    u = (q & ((1 << bits) - 1)).reshape(-1)
    stream = ((u[:, None] >> np.arange(bits)) & 1).astype(np.uint8).reshape(-1)
    return np.packbits(stream, bitorder="little").tobytes() + s.astype(np.float32).tobytes()


BAND = ("hconv/gate_conv", "hconv/value_conv", "hconv/out_conv")
FP32 = {"dense/out/kernel", "dense/bias", "dense/protos/kernel", "dense/coef/kernel", "reader/out/bias"}


def band_taps(rows, k):
    """[rows, k] mask of the taps a band applies to: tap t of output j reads channel j + t - (k - 1) / 2."""
    j, t = np.arange(rows)[:, None], np.arange(k)[None, :]
    c = j + t - (k - 1) // 2
    return (c >= 0) & (c < rows)


def cmd_export(argv):
    """Write an .oasis weight file for the Oasis engine.

Layout: 8-byte magic "OASIS\\0v1", u32 header length, UTF-8 header of "key value" lines, u32 tensor count,
then per tensor: u16 name length, name, u8 dtype, u8 bits, u8 ndim, u8 flags, u32 shape[ndim], u64 offset,
u64 nbytes; blobs follow. Flag 1 marks a deflated blob (u64 length of the stored bytes, then a zlib stream of
them); other blobs are 64-byte aligned. Dtypes: 0 f32, 1 f16, 2 cq, 3 raw bytes, 4 integer rows (signed
integers of `bits` bits in one LSB-first bit stream, then one f32 scale per row), 5 cq with 8-bit scales. A cq
tensor is [out, in]: every row's packed indices (LSB first, columns padded to 128), then one scale per row and
128-column group, as fp16 (dtype 2) or as an 8-bit code (dtype 5) followed by two f32 values lo and hi: code 0
is zero and codes 1..255 step evenly in log scale from exp(lo) to exp(hi). --band-bits stores the HConvGLU
bands as integer rows, with the taps that fall outside the channel vector, which the engine never reads,
stored as zero; --scale-bits 8 stores the group scales of every quantised matrix outside --keep-scales as
codes.

Quantised tensors whose latent weights match Mirage's are copied from mirage.cact; others are quantised here
with the training-time quantiser, on the device the weights were trained on (--device, CUDA when available):
quantisation-aware training leaves many latent weights at codebook boundaries, where the backend's float32
rounding decides the index."""
    ap = parser("export", cmd_export.__doc__)
    ap.add_argument("--weights", help="float32 weights to export (default: Mirage's)")
    ap.add_argument("--base", help="Mirage's float32 weights (default: from --mirage)")
    ap.add_argument("--archive", help="Mirage's archive (default: from --mirage)")
    ap.add_argument("--out", required=True)
    ap.add_argument("--name", default="oasis")
    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    ap.add_argument("--band-bits", type=int, choices=[8, 10, 12, 16], default=12)
    ap.add_argument("--scale-bits", type=int, choices=[8, 16], default=8)
    ap.add_argument("--keep-scales", default="reader/", help="comma list of name prefixes whose group scales stay fp16")
    ap.add_argument("--no-deflate", action="store_true", help="store every tensor as it is, without zlib")
    args = ap.parse_args(argv)
    keep_scales = tuple(p for p in args.keep_scales.split(",") if p)
    arc = Archive(args.archive or mirage_archive(args))
    w = {k: v.float() for k, v in load_file(args.weights or mirage_weights(args)).items() if not k.startswith(("mtp_", "name_head"))}
    base = load_file(args.base or mirage_weights(args))
    vocab = vocab_of(arc)
    quantizer = Quantizer(arc.codebook, args.device)

    tensors = []   # (name, dtype, bits, shape, bytes)
    quantised = {}
    for item in cq_plan():
        key = item["name"] if item["layer"] is None else f"{item['name']}/{item['layer']}"
        src = w[item["name"]] if item["layer"] is None else w[item["name"]][item["layer"]]
        ref_src = base[item["name"]] if item["layer"] is None else base[item["name"]][item["layer"]]
        e = arc.entries[item["entry"]]
        if torch.equal(src.float(), ref_src.float()):
            blob = arc.raw(item["entry"])
        else:
            R = plan_get({item["name"]: w[item["name"]]}, item).float()
            blob = cq_blob_from_latent(R, item["bits"], quantizer)
        quantised.setdefault(item["name"], []).append(item["layer"])
        rows, cols = int(np.prod(e.shape[:-1])), int(e.shape[-1])
        if args.scale_bits == 8 and not item["name"].startswith(keep_scales):
            tensors.append((key, 5, e.bits, list(e.shape), compact_scales(blob, rows, cols, e.bits)))
        else:
            tensors.append((key, 2, e.bits, list(e.shape), blob))
    for k, v in sorted(w.items()):
        if k in quantised:
            continue
        if k == "embedding/embedding":
            continue
        if k in FP32:
            tensors.append((k, 0, 0, list(v.shape), v.float().numpy().tobytes()))
            continue
        if args.band_bits < 16 and k.endswith(BAND):
            keep = band_taps(*v.shape)
            tensors.append((k, 4, args.band_bits, list(v.shape), int_rows(v.numpy(), args.band_bits, keep)))
            continue
        if k.startswith("stack/layers/block/") and v.dim() >= 1 and v.shape[0] == 8:
            for l in range(8):
                tensors.append((f"{k}/{l}", 1, 0, list(v[l].shape) or [1], v[l].half().numpy().tobytes()))
        else:
            tensors.append((k, 1, 0, list(v.shape) or [1], v.half().numpy().tobytes()))
    tensors.append(("hmlp/perm1", 0, 0, [512], np.random.RandomState(11).permutation(512).astype(np.float32).tobytes()))
    tensors.append(("hmlp/perm2", 0, 0, [512], np.random.RandomState(13).permutation(512).astype(np.float32).tobytes()))
    tensors.append(("quant/codebook", 0, 0, [len(arc.codebook)], np.asarray(arc.codebook, np.float32).tobytes()))
    charset = arc.array(CHARSET_ENTRY).astype(np.float32)
    tensors.append(("reader/charset", 0, 0, [len(charset)], charset.tobytes()))
    pieces = bytearray(struct.pack("<I", vocab.n_pieces))
    for p, t in zip(vocab.pieces, vocab.types):
        b = p.encode("utf-8")
        pieces += struct.pack("<BH", t, len(b)) + b
    tensors.append(("tokenizer/pieces", 3, 0, [len(pieces)], bytes(pieces)))

    langs = " ".join(f"{k}:{v}" for k, v in vocab.lang_ids.items())
    # The vocabulary carries eleven language tags; Mirage was trained to describe in these seven.
    languages = "en de fr es it nl pl"
    header = "\n".join([
        f"name {args.name}", "format 1", "d_model 512", f"vocab {vocab.size}", f"n_pieces {vocab.n_pieces}",
        f"bos {vocab.bos}", f"eos {vocab.eos}", f"pad {vocab.pad}", f"ref_base {vocab.ref_base}", "n_obj_refs 64",
        "n_txt_refs 32", f"langs {langs}", f"languages {languages}", "group 128", "canvas 512", "engram_sites 3 7",
        "engram_slots 18432", "rope_dec 100000", "rope_enc 10000", "codebook_bits 2:0 3:4 4:12",
    ]).encode("utf-8")

    # A tensor whose bytes deflate by more than half a percent is stored as flag 1: its u64 length, then the
    # zlib stream; the engine inflates it at load. Stored tensors start on 64-byte boundaries, deflated ones
    # anywhere.
    stored = []
    for name, dt, bits, shape, blob in tensors:
        flags = 0
        if not args.no_deflate:
            z = zlib.compress(bytes(blob), 9)
            if len(z) + 8 < len(blob) * 0.995:
                blob, flags = struct.pack("<Q", len(blob)) + z, 1
        stored.append((name, dt, bits, shape, blob, flags))
    offsets, pos = [], 0
    for name, dt, bits, shape, blob, flags in stored:
        if not flags:
            pos = -(-pos // 64) * 64
        offsets.append(pos)
        pos += len(blob)
    head = bytearray(b"OASIS\0v1") + struct.pack("<I", len(header)) + header + struct.pack("<I", len(stored))
    for (name, dt, bits, shape, blob, flags), off in zip(stored, offsets):
        nb = name.encode("utf-8")
        head += struct.pack("<H", len(nb)) + nb + struct.pack("<BBBB", dt, bits, len(shape), flags)
        head += struct.pack(f"<{len(shape)}I", *shape) + struct.pack("<QQ", off, len(blob))
    data_start = -(-len(head) // 64) * 64
    with open(args.out, "wb") as f:
        f.write(head + b"\0" * (data_start - len(head)))
        for (name, dt, bits, shape, blob, flags), off in zip(stored, offsets):
            f.seek(data_start + off)
            f.write(blob)
    size = os.path.getsize(args.out)
    n_q = sum(1 for t in tensors if t[1] in (2, 5))
    print(f"wrote {args.out}: {len(tensors)} tensors ({n_q} quantised on {args.device}), {size / 1e6:.2f} MB, data at {data_start}")


def cmd_checkpoint(argv):
    """Assemble an Oasis float32 checkpoint: Mirage's weights with replacement components, such as a retrained
reader. Training-only heads (multi-token prediction, name head) are dropped. The metadata records the base
checkpoint's hash and every replacement file's hash."""
    ap = parser("checkpoint", cmd_checkpoint.__doc__)
    ap.add_argument("--base", help="Mirage's float32 weights (default: from --mirage)")
    ap.add_argument("--replace", action="append", default=[], help="a safetensors file of replacement tensors; repeatable")
    ap.add_argument("--dead-offset", type=float, default=0.0, help="added to the reader biases of classes the base never trained")
    ap.add_argument("--out", required=True)
    args = ap.parse_args(argv)
    base_path = args.base or mirage_weights(args)
    w = {k: v.float() for k, v in load_file(base_path).items() if not k.startswith(("mtp_", "name_head"))}
    dead = dead_classes(w)
    meta = {"base": os.path.basename(base_path), "base_sha256": sha256(base_path), "replacements": []}
    for path in args.replace:
        r = load_file(path)
        unknown = [k for k in r if k not in w]
        if unknown:
            raise SystemExit(f"{path} has tensors the base lacks: {unknown[:5]}")
        w.update({k: v.float() for k, v in r.items()})
        meta["replacements"].append({"file": os.path.basename(path), "sha256": sha256(path), "tensors": len(r)})
    if args.dead_offset:
        w["reader/out/bias"][dead] += args.dead_offset
        meta["reader_dead_class_bias_offset"] = args.dead_offset
    os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
    save_file({k: v.contiguous() for k, v in w.items()}, args.out, metadata={"oasis": json.dumps(meta)})
    print(json.dumps(meta, indent=1))
    print(f"wrote {args.out}: {len(w)} tensors, {sum(v.numel() for v in w.values()):,} parameters")


def cmd_run(argv):
    """Run the reference forward pass on one image: detections, OCR, masks and a greedy decoder probe. Writes
annotated.png, the two heat maps and result.json to --out (default: runs/IMAGE_NAME)."""
    ap = parser("run", cmd_run.__doc__)
    ap.add_argument("image")
    ap.add_argument("--weights", help="float32 weights (default: Mirage's)")
    ap.add_argument("--archive", help="Mirage's archive (default: from --mirage)")
    ap.add_argument("--prefix", default="<s>", help="decoder prompt: comma list of pieces or token ids")
    ap.add_argument("--max-new", type=int, default=200)
    ap.add_argument("--conf", type=float, default=0.3)
    ap.add_argument("--text-conf", type=float, default=0.3)
    ap.add_argument("--out", default=None)
    ap.add_argument("--set", action="append", default=[], help="OPTS override key=value")
    args = ap.parse_args(argv)
    for kv in args.set:
        k, v = kv.split("=", 1)
        OPTS[k] = type(OPTS[k])(eval(v, {}, {})) if not isinstance(OPTS[k], str) else v

    arc = Archive(args.archive or mirage_archive(args))
    vocab = vocab_of(arc)
    charset = charset_of(arc)
    t0 = time.time()
    model = Mirage(args.weights or mirage_weights(args))
    img = Image.open(args.image).convert("RGB")
    enc = model.encode(img)
    t1 = time.time()
    det = model.detect(enc, conf=args.conf, text_conf=args.text_conf, nms_iou=0.5)
    meta = enc["meta"]
    print(f"encode {t1 - t0:.2f}s  image {meta['W']}x{meta['H']}  canvas valid {meta['w']}x{meta['h']}")
    for kind in ("object", "text"):
        print(f"-- {kind}: {len(det[kind])}")
        for k, b in enumerate(det[kind][:20]):
            print(f"   {k:2d} score {b['score']:.3f} cell {b['cell']} box1000 {model.to_1000(b['box'], meta)}")

    ocr = []
    for k, b in enumerate(det["text"]):
        bb = model.to_1000(b["box"], meta)
        px = (bb[0] / 1000 * meta["W"], bb[1] / 1000 * meta["H"], bb[2] / 1000 * meta["W"], bb[3] / 1000 * meta["H"])
        labels, _ = model.read(img, px)
        ocr.append(ctc_text(labels, charset))
        print(f"   ocr {k:2d}: {ocr[-1]!r}")

    masks = model.masks(enc, det["object"][:8])

    out_dir = args.out or os.path.join(HERE, "runs", os.path.splitext(os.path.basename(args.image))[0])
    os.makedirs(out_dir, exist_ok=True)
    vis = img.copy()
    dr = ImageDraw.Draw(vis, "RGBA")
    for k, (b, m) in enumerate(zip(det["object"][:8], masks)):
        mm = (m.numpy() > 0.5).astype(np.uint8) * 255
        sc = CANVAS / meta["r"] / 128
        mimg = Image.fromarray(mm).resize((int(128 * sc), int(128 * sc)), Image.NEAREST).crop((0, 0, meta["W"], meta["H"]))
        col = [(230, 25, 75), (60, 180, 75), (0, 130, 200), (245, 130, 48), (145, 30, 180), (70, 240, 240), (240, 50, 230), (210, 245, 60)][k]
        vis.paste(Image.new("RGB", vis.size, col), (0, 0), mimg.point(lambda v: 90 if v else 0))
    for kind, color in (("object", (255, 0, 0, 255)), ("text", (255, 205, 0, 255))):
        for k, b in enumerate(det[kind]):
            bb = model.to_1000(b["box"], meta)
            xy = [bb[0] / 1000 * meta["W"], bb[1] / 1000 * meta["H"], bb[2] / 1000 * meta["W"], bb[3] / 1000 * meta["H"]]
            dr.rectangle(xy, outline=color, width=2)
            dr.text((xy[0] + 2, xy[1] + 2), f"{kind[0]}{k} {b['score']:.2f}", fill=color)
    vis.save(os.path.join(out_dir, "annotated.png"))
    heat = torch.sigmoid(det["raw"][..., :2]).numpy()
    for ch, name in ((0, "heat0"), (1, "heat1")):
        Image.fromarray((heat[..., ch] * 255).astype(np.uint8)).resize((256, 256), Image.NEAREST).save(os.path.join(out_dir, f"{name}.png"))

    M, mvalid, refs = model.memory(enc, det["object"], det["text"], vocab)
    dec = Decoder(model, M, mvalid)
    prefix = []
    for item in args.prefix.split(","):
        item = item.strip()
        if not item:
            continue
        prefix.append(int(item) if item.isdigit() else vocab.id_of[item])
    banned = torch.zeros(vocab.size, dtype=torch.bool)
    for i in vocab.lang_ids.values():
        banned[i] = True
    banned[vocab.ref_base:] = True
    for r in refs:
        banned[r] = False
    t2 = time.time()
    logits = None
    for tid in prefix:
        logits = dec.step(tid)
    gen = []
    for _ in range(args.max_new):
        lg = logits.masked_fill(banned, float("-inf"))
        nxt = int(lg.argmax())
        gen.append(nxt)
        if nxt == vocab.eos:
            break
        logits = dec.step(nxt)
    t3 = time.time()
    text = vocab.decode(gen)
    print(f"decoder: prefix {prefix} -> {len(gen)} tokens in {t3 - t2:.1f}s")
    print(text)
    json.dump(dict(meta=meta, objects=[dict(box=model.to_1000(b['box'], meta), score=b['score']) for b in det['object']],
                   text=[dict(box=model.to_1000(b['box'], meta), score=b['score'], ocr=o) for b, o in zip(det['text'], ocr)],
                   prefix=prefix, generated=gen, decoded=text),
              open(os.path.join(out_dir, "result.json"), "w", encoding="utf-8"), ensure_ascii=False, indent=1)


# ===================================================================================== commands: reader
class RealCrops(torch.utils.data.IterableDataset):
    """Real word crops (parquet rows of image bytes and text) as reader strips with CTC labels; each worker
    reads its own share of the row groups. jitter: random box margins and rescaling blur."""

    def __init__(self, files, index, jitter):
        super().__init__()
        self.files, self.index, self.jitter = files, index, jitter

    def __iter__(self):
        import pyarrow.parquet as pq
        info = torch.utils.data.get_worker_info()
        wid, nw = (info.id, info.num_workers) if info else (0, 1)
        rng = random.Random(4242 + wid)
        units = [(f, g) for f in self.files for g in range(pq.ParquetFile(f).num_row_groups)]
        mine = [u for i, u in enumerate(units) if i % nw == wid]
        while True:
            rng.shuffle(mine)
            for f, g in mine:
                rows = pq.ParquetFile(f).read_row_group(g).to_pylist()
                rng.shuffle(rows)
                for row in rows:
                    text = row["text"]
                    if not text or text == "###" or any(c not in self.index for c in text):
                        continue
                    try:
                        img = Image.open(io.BytesIO(row["image"]["bytes"])).convert("RGB")
                    except Exception:
                        continue
                    w, h = img.size
                    if h < 6 or w < 4:
                        continue
                    if self.jitter:
                        a = np.asarray(img, dtype=np.float32)
                        fill = tuple(int(v) for v in np.concatenate([a[0], a[-1], a[:, 0], a[:, -1]]).mean(0))
                        side = [int(rng.uniform(-0.06, 0.3) * h) for _ in range(4)]
                        canvas = Image.new("RGB", (w + max(0, side[0]) + max(0, side[2]), h + max(0, side[1]) + max(0, side[3])), fill)
                        canvas.paste(img, (max(0, side[0]), max(0, side[1])))
                        x0, y0 = max(0, -side[0]), max(0, -side[1])
                        x1, y1 = canvas.width - max(0, -side[2]), canvas.height - max(0, -side[3])
                        if x1 - x0 >= 4 and y1 - y0 >= 6:
                            img = canvas.crop((x0, y0, x1, y1))
                        if rng.random() < 0.5:
                            f = rng.uniform(0.5, 1.0)
                            img = img.resize((max(1, int(img.width * f)), max(1, int(img.height * f))), Image.BILINEAR)
                        w, h = img.size
                    U = max(1, int(round(48 * w / h)))
                    if U > 320:
                        continue
                    labels = [self.index[c] for c in text]
                    reps = sum(1 for a, b in zip(labels, labels[1:]) if a == b)
                    if -(-U // 4) < len(labels) + reps:
                        continue
                    strip = np.full((48, 320, 3), 127, np.uint8)
                    strip[:, :U] = np.asarray(img.resize((U, 48), Image.BILINEAR))
                    yield strip, U, labels, text


class SyntheticLines(torch.utils.data.IterableDataset):
    """Synthetic text lines from Generator, seeded per worker."""

    def __init__(self, charset, dead, words, fonts, style, tight):
        super().__init__()
        self.charset, self.dead, self.words, self.fonts, self.style, self.tight = charset, dead, words, fonts, style, tight

    def __iter__(self):
        info = torch.utils.data.get_worker_info()
        gen = Generator(self.charset, self.dead, self.words, self.fonts, 9000 + (info.id if info else 0) + random.randint(0, 10**6),
                        style=self.style, tight=self.tight)
        while True:
            yield gen.sample()


def collate_strips(batch):
    x = torch.from_numpy(np.stack([b[0] for b in batch])).permute(0, 3, 1, 2).float() / 127.5 - 1.0
    lens = torch.tensor([-(-b[1] // 4) for b in batch])
    tgt = torch.tensor([c for b in batch for c in b[2]])
    tl = torch.tensor([len(b[2]) for b in batch])
    return x, lens, tgt, tl


def batches(ds, batch, workers):
    return iter(torch.utils.data.DataLoader(ds, batch_size=batch, num_workers=workers, collate_fn=collate_strips,
                                            prefetch_factor=4, persistent_workers=True))


def cmd_train(argv):
    """Quantisation-aware fine-tuning of the CTC reader on real word crops mixed with synthetic lines.

The 2-bit matrices run through the stock quantiser in every forward pass (straight-through gradients), so the
exported 2-bit reader is the trained one. Real crops get CTC plus distillation toward Mirage's 2-bit reader;
synthetic lines carrying the untrained characters get CTC. Checkpoints of the reader's tensors and a log are
written to --out; `checkpoint --replace` puts one into the full weights.

Oasis's reader: --real 'data/ocr/textocr/train-*.parquet' --act-quant --jitter for 8000 steps, then
--init on the result with --synth realistic for 4000, then again with --tight 0.5 for 3000."""
    ap = parser("train", cmd_train.__doc__)
    ap.add_argument("--weights", help="Mirage's float32 weights (default: from --mirage)")
    ap.add_argument("--init", default=None, help="reader weights to start from instead of the original")
    ap.add_argument("--real", required=True, help="glob of parquet files of word crops (image bytes and text)")
    ap.add_argument("--evalsets", help="comma list of packed crop archives (default: data/ocr/real_words.npz and synth_val.npz)")
    ap.add_argument("--steps", type=int, default=8000)
    ap.add_argument("--real-batch", type=int, default=128)
    ap.add_argument("--syn-batch", type=int, default=64)
    ap.add_argument("--syn-weight", type=float, default=1.0)
    ap.add_argument("--real-weight", type=float, default=1.0, help="weight of CTC on real labels (0: distillation only)")
    ap.add_argument("--distill", type=float, default=0.5)
    ap.add_argument("--lr", type=float, default=1e-4)
    ap.add_argument("--act-quant", action="store_true", help="engine-style int8 activations in student and teacher")
    ap.add_argument("--jitter", action="store_true", help="random box margins and rescaling blur on real crops")
    ap.add_argument("--synth", choices=["uniform", "realistic"], default="uniform", help="synthetic text distribution")
    ap.add_argument("--fonts", default=",".join(LINUX_FONTS), help="comma list of font globs for synthetic lines")
    ap.add_argument("--tight", type=float, default=0.0, help="share of synthetic lines cut as tightly as a detector box")
    ap.add_argument("--freeze", default="", help="comma list of name prefixes to freeze, e.g. w0,dw_0,pw_0")
    ap.add_argument("--workers", type=int, default=12)
    ap.add_argument("--eval-every", type=int, default=1000)
    ap.add_argument("--out", required=True)
    args = ap.parse_args(argv)
    os.makedirs(args.out, exist_ok=True)
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    torch.manual_seed(0)

    arc = Archive(mirage_archive(args))
    charset = charset_of(arc)
    index = {c: i + 1 for i, c in enumerate(charset)}
    words = vocab_words(vocab_of(arc))
    base = load_file(args.weights or mirage_weights(args))
    dead_cls = dead_classes(base)
    dead_chars = [charset[c - 1] for c in dead_cls]
    start = dict(base)
    if args.init:
        start.update(load_file(args.init))
    quantizer = Quantizer(arc.codebook, dev)
    model = Reader(start, quantizer, act_quant=args.act_quant).to(dev)
    teacher = Reader(base, quantizer, act_quant=args.act_quant).to(dev).eval()
    if not args.init:
        # Untrained classes start from the output weights of the letter they are closest to.
        K = model.params["reader__out__kernel"]
        Bv = model.params["reader__out__bias"]
        special = {"–": "-", "¡": "!", "ĸ": "k", "ʼn": "n", "Ŋ": "N", "ŋ": "n", "Ŧ": "T", "ŧ": "t", "Ð": "D",
                   "Þ": "P", "þ": "p", "Ŀ": "L", "ŀ": "l", "IJ": "I", "ij": "i"}
        with torch.no_grad():
            for c in dead_cls:
                ch = charset[c - 1]
                cands = ([special[ch]] if ch in special else []) + [v for v in charset if v != ch and v not in dead_chars and unicodedata.normalize("NFKD", v)[0] == unicodedata.normalize("NFKD", ch)[0]] + [unicodedata.normalize("NFKD", ch)[0]]
                d = next((charset.index(v) + 1 for v in cands if v in charset and v not in dead_chars), None)
                if d is not None:
                    K[:, c] = K[:, d]
                    Bv[c] = Bv[d] - 0.5
    for q in teacher.parameters():
        q.requires_grad_(False)
    frozen = [f for f in args.freeze.split(",") if f]
    trainable = []
    for name in model.names:
        q = model.params[name.replace("/", "__")]
        short = name[len("reader/"):]
        keep = not any(short == f or short.startswith(f + "/") for f in frozen)
        q.requires_grad_(keep)
        if keep:
            trainable.append(q)
    print(f"training {sum(q.numel() for q in trainable):,} of {sum(q.numel() for q in model.parameters()):,} reader weights", flush=True)
    real_files = sorted(glob.glob(args.real))
    print("real files:", real_files, flush=True)

    evalsets = args.evalsets.split(",") if args.evalsets else [os.path.join(args.data, "ocr", "real_words.npz"),
                                                               os.path.join(args.data, "ocr", "synth_val.npz")]
    sets = [load_packed(p) for p in evalsets]
    dset = set(dead_chars)
    res = evaluate_reader(teacher, sets, charset, dset, dev)
    print("original:", json.dumps(res), flush=True)
    log = open(os.path.join(args.out, "log.jsonl"), "a")
    log.write(json.dumps(dict(step=0, eval=res, args=vars(args))) + "\n")

    real_it = batches(RealCrops(real_files, index, args.jitter), args.real_batch, max(1, args.workers * 2 // 3))
    syn_it = batches(SyntheticLines(charset, dead_chars, words, args.fonts.split(","), args.synth, args.tight),
                     args.syn_batch, max(1, args.workers // 3))
    opt = torch.optim.AdamW(trainable, lr=args.lr, weight_decay=0.0)
    sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1.0, (s + 1) / 200) * 0.5 * (1 + math.cos(math.pi * min(s, args.steps) / args.steps)))
    ctc = torch.nn.CTCLoss(blank=0, zero_infinity=True)
    model.train()
    t0 = time.time()
    for step in range(1, args.steps + 1):
        xr, lr_, tr, tlr = next(real_it)
        xs, ls_, ts, tls = next(syn_it)
        xr, xs = xr.to(dev), xs.to(dev)
        lp_r = F.log_softmax(model(xr), -1)
        lp_s = F.log_softmax(model(xs), -1)
        l_real = ctc(lp_r.transpose(0, 1), tr.to(dev), lr_.to(dev), tlr.to(dev))
        l_syn = ctc(lp_s.transpose(0, 1), ts.to(dev), ls_.to(dev), tls.to(dev))
        with torch.no_grad():
            tp = F.softmax(teacher(xr), -1)
        kd = F.kl_div(lp_r, tp, reduction="batchmean") / lp_r.shape[1]
        loss = args.real_weight * l_real + args.syn_weight * l_syn + args.distill * kd
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(trainable, 5.0)
        opt.step()
        sched.step()
        if step % 100 == 0:
            print(f"step {step} real {l_real.item():.3f} syn {l_syn.item():.3f} kd {kd.item():.4f} lr {sched.get_last_lr()[0]:.2e} {time.time() - t0:.0f}s", flush=True)
        if step % args.eval_every == 0 or step == args.steps:
            res = evaluate_reader(model, sets, charset, dset, dev)
            print(f"eval step {step}: {json.dumps(res)}", flush=True)
            log.write(json.dumps(dict(step=step, eval=res)) + "\n")
            log.flush()
            save_file(model.state(), os.path.join(args.out, f"reader_{step:06d}.safetensors"))
    print("done", flush=True)


def reader_name(r):
    return "base" if r == "base" else os.path.basename(os.path.dirname(r)) + "/" + os.path.basename(r).replace(".safetensors", "")


def cmd_eval_reader(argv):
    """Batched reader evaluation on packed crops, with the engine's window and frame-retention rules.

Reports, per source: word accuracy (lowercase, letters and digits), exact match, CER, and recall of the
characters the original checkpoint never trained. --dead-offset adds a constant to those classes' biases,
which trades their recall against false readings on ordinary text."""
    ap = parser("eval-reader", cmd_eval_reader.__doc__)
    ap.add_argument("packed", nargs="+")
    ap.add_argument("--reader", action="append", default=[], help="reader weights to score in place of Mirage's; repeatable")
    ap.add_argument("--dead-offset", type=float, action="append", default=[])
    ap.add_argument("--weights", help="Mirage's float32 weights (default: from --mirage)")
    ap.add_argument("--qat", action="store_true", help="run the 2-bit matrices through the training-time quantiser")
    ap.add_argument("--act-quant", action="store_true", help="int8 activations per rotated 128-group, as the engine computes them")
    args = ap.parse_args(argv)
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    arc = Archive(mirage_archive(args))
    charset = charset_of(arc)
    base = load_file(args.weights or mirage_weights(args))
    dead_cls = dead_classes(base)
    dead_chars = set(charset[c - 1] for c in dead_cls)
    sets = [load_packed(f) for f in args.packed]
    for r in (args.reader or ["base"]):
        w = dict(base)
        if r != "base":
            w.update(load_file(r))
        net = Reader(w, Quantizer(arc.codebook, dev) if args.qat else None, act_quant=args.act_quant).to(dev).eval()
        b0 = net.params["reader__out__bias"].detach().clone()
        for off in (args.dead_offset or [0.0]):
            with torch.no_grad():
                net.params["reader__out__bias"].copy_(b0)
                net.params["reader__out__bias"][dead_cls] += off
            for X, owner, keeps, labels, sources in sets:
                stats = {}
                for i, path in enumerate(reader_paths(net, X, owner, keeps, len(labels), dev)):
                    pred, gt = ctc_text(path, charset), str(labels[i])
                    st = stats.setdefault(str(sources[i]), dict(n=0, acc=0, exact=0, ed=0, len=0, dh=0, dt=0, fp=0))
                    st["n"] += 1
                    st["acc"] += word_key(pred) == word_key(gt)
                    st["exact"] += pred == gt
                    st["ed"] += edit_distance(pred, gt)
                    st["len"] += len(gt)
                    for ch in set(gt) & dead_chars:
                        st["dt"] += gt.count(ch)
                        st["dh"] += min(gt.count(ch), pred.count(ch))
                    st["fp"] += sum(1 for ch in pred if ch in dead_chars and ch not in gt)
                name = reader_name(r)
                for src, st in sorted(stats.items()):
                    n = st["n"]
                    print(f"{src:20s} {name:26s} off={off:+6.2f} n={n} word_acc={st['acc'] / n:.4f} exact={st['exact'] / n:.4f} "
                          f"cer={st['ed'] / max(st['len'], 1):.4f} dead_recall={st['dh'] / max(st['dt'], 1):.3f} dead_false={st['fp']}", flush=True)


def cmd_eval_detected(argv):
    """Score readers on packed detector crops (pack-detected): the crops of a word read and joined left to
right, against the word."""
    ap = parser("eval-detected", cmd_eval_detected.__doc__)
    ap.add_argument("packed", nargs="+")
    ap.add_argument("--reader", action="append", default=[], help="reader weights to score in place of Mirage's; repeatable")
    ap.add_argument("--qat", action="store_true", help="run the 2-bit matrices through the training-time quantiser")
    ap.add_argument("--act-quant", action="store_true", help="int8 activations per rotated 128-group, as the engine computes them")
    ap.add_argument("--dead-offset", type=float, action="append", default=[])
    ap.add_argument("--weights", help="Mirage's float32 weights (default: from --mirage)")
    args = ap.parse_args(argv)
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    arc = Archive(mirage_archive(args))
    charset = charset_of(arc)
    base = load_file(args.weights or mirage_weights(args))
    dead_cls = dead_classes(base)
    dead = set(charset[c - 1] for c in dead_cls)

    sets = []
    for f in args.packed:
        d = dict(np.load(f, allow_pickle=False))
        pix, offs = d["pix"], d["offsets"]
        wins, owner, keeps = [], [], []
        for i in range(len(offs) - 1):
            ws, ks = reader_windows(pix[:, offs[i]:offs[i + 1]])
            wins += ws
            owner += [i] * len(ws)
            keeps += ks
        X = torch.from_numpy(np.stack(wins)).permute(0, 3, 1, 2) if wins else torch.zeros(0, 3, 48, 320, dtype=torch.uint8)
        sets.append((os.path.basename(f), X, owner, keeps, d["item"], d["x0"], d["labels"], d["sources"]))

    for r in (args.reader or ["base"]):
        w = dict(base)
        if r != "base":
            w.update(load_file(r))
        q = Quantizer(arc.codebook, dev) if args.qat else None
        net = Reader(w, q, act_quant=args.act_quant).to(dev).eval()
        b0 = net.params["reader__out__bias"].detach().clone()
        for off in (args.dead_offset or [0.0]):
            with torch.no_grad():
                net.params["reader__out__bias"].copy_(b0)
                net.params["reader__out__bias"][dead_cls] += off
            for name, X, owner, keeps, item, x0, labels, sources in sets:
                crop_text = [ctc_text(path, charset) for path in reader_paths(net, X, owner, keeps, len(item), dev)]
                words = [[] for _ in range(len(labels))]
                for c, (wi, xx) in enumerate(zip(item, x0)):
                    words[wi].append((xx, crop_text[c]))
                stats = {}
                for wi, gt in enumerate(labels):
                    pred = " ".join(t for _, t in sorted(words[wi])).strip()
                    st = stats.setdefault(str(sources[wi]), [0, 0, 0, 0, 0])
                    st[0] += 1
                    st[1] += word_key(pred) == word_key(str(gt))
                    st[2] += edit_distance(pred, str(gt))
                    st[3] += len(str(gt))
                    st[4] += sum(1 for ch in pred if ch in dead and ch not in str(gt))
                nm = reader_name(r)
                for src, (n, a, e, ln, fp) in sorted(stats.items()):
                    print(f"{src:28s} {nm:26s} off={off:+5.2f} n={n} word_acc={a / n:.4f} cer={e / max(ln, 1):.4f} dead_false={fp}", flush=True)


def cmd_pack_words(argv):
    """Pack cropped-word OCR benchmarks (parquet files of image bytes and text) into a numpy archive, each crop
resized to height 48 with its aspect kept: pix uint8 [48, total_width, 3], offsets int64 [N+1], labels and
sources."""
    ap = parser("pack-words", cmd_pack_words.__doc__)
    ap.add_argument("out")
    ap.add_argument("files", nargs="+")
    ap.add_argument("--limit", type=int, default=0, help="the first N rows of each file")
    args = ap.parse_args(argv)
    import pyarrow.parquet as pq
    strips, labels, sources = [], [], []
    for f in args.files:
        table = pq.read_table(f)
        if args.limit:
            table = table.slice(0, args.limit)
        for row in table.to_pylist():
            img = Image.open(io.BytesIO(row["image"]["bytes"])).convert("RGB")
            w, h = img.size
            U = max(1, int(round(48 * w / h)))
            strips.append(np.asarray(img.resize((U, 48), Image.BILINEAR), dtype=np.uint8))
            labels.append(row["text"])
            sources.append(os.path.basename(f))
    offsets = np.cumsum([0] + [s.shape[1] for s in strips]).astype(np.int64)
    np.savez_compressed(args.out, pix=np.concatenate(strips, 1), offsets=offsets, labels=np.array(labels), sources=np.array(sources))
    print(len(labels), "crops, total width", int(offsets[-1]))


def cmd_pack_detected(argv):
    """Pack the crops the Oasis engine's text detector cuts from word images placed in scenes as bench-ocr
places them: each text box is cropped from the scene with the engine's rounding and resized to 48 px height,
so that readers can be scored on exactly the crops the engine reads. Arrays: pix [48, total, 3], offsets, item
(crop -> word index), x0 (left edge, for ordering), labels (per word) and sources (per word)."""
    ap = parser("pack-detected", cmd_pack_detected.__doc__)
    ap.add_argument("out")
    ap.add_argument("files", nargs="+")
    ap.add_argument("--oasis-lib", help="the Oasis engine library (default: found as oasis.py finds it)")
    ap.add_argument("--oasis-model", required=True)
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--text-conf", type=float, default=0.2)
    args = ap.parse_args(argv)
    import pyarrow.parquet as pq
    eng = OasisEngine(args.oasis_lib, args.oasis_model)
    opts = {"sections": ["text"], "text_conf": args.text_conf, "max_text_lines": 4, "max_objects": 1, "max_tokens": 1, "mask_format": "largest"}
    strips, item, xs, labels, sources = [], [], [], [], []
    for f in args.files:
        t = pq.read_table(f)
        if args.limit:
            t = t.slice(0, args.limit)
        for r in t.to_pylist():
            scene = word_scene(Image.open(io.BytesIO(r["image"]["bytes"])).convert("RGB"))
            rec, _, _ = eng.describe(np.asarray(scene, dtype=np.uint8).tobytes(), 640, 480, opts)
            wi = len(labels)
            labels.append(r["text"])
            sources.append(os.path.basename(f))
            for tbox in rec.get("text", []) if rec else []:
                b = tbox["bbox"]
                x0, y0, x1, y1 = b[0] / 1000 * 640, b[1] / 1000 * 480, b[2] / 1000 * 640, b[3] / 1000 * 480
                crop = scene.crop((math.floor(x0), math.floor(y0), math.ceil(x1), math.ceil(y1)))
                if crop.width < 1 or crop.height < 1:
                    continue
                U = max(1, int(round(48 * crop.width / crop.height)))
                strips.append(np.asarray(crop.resize((U, 48), Image.BILINEAR), dtype=np.uint8))
                item.append(wi)
                xs.append(b[0])
    offsets = np.cumsum([0] + [s.shape[1] for s in strips]).astype(np.int64)
    np.savez_compressed(args.out, pix=np.concatenate(strips, 1) if strips else np.zeros((48, 0, 3), np.uint8), offsets=offsets,
                        item=np.array(item, np.int64), x0=np.array(xs, np.float32), labels=np.array(labels), sources=np.array(sources))
    print(f"{len(labels)} words, {len(strips)} detected crops")


def cmd_synth_val(argv):
    """Write a fixed synthetic validation set in the packed crop format, with the characters Mirage's reader
never learned listed under "dead"."""
    ap = parser("synth-val", cmd_synth_val.__doc__)
    ap.add_argument("out")
    ap.add_argument("--n", type=int, default=2000)
    ap.add_argument("--seed", type=int, default=12345)
    ap.add_argument("--style", choices=["uniform", "realistic"], default="uniform")
    ap.add_argument("--tight", type=float, default=0.0)
    ap.add_argument("--fonts", default=",".join(LINUX_FONTS), help="comma list of font globs")
    args = ap.parse_args(argv)
    arc = Archive(mirage_archive(args))
    charset = charset_of(arc)
    words = vocab_words(vocab_of(arc))
    dead = [charset[c - 1] for c in dead_classes(load_file(mirage_weights(args)))]
    gen = Generator(charset, dead, words, args.fonts.split(","), args.seed, style=args.style, tight=args.tight)
    strips, labels = [], []
    for _ in range(args.n):
        strip, U, _, s = gen.sample()
        strips.append(strip[:, :U])
        labels.append(s)
    offsets = np.cumsum([0] + [s.shape[1] for s in strips]).astype(np.int64)
    np.savez_compressed(args.out, pix=np.concatenate(strips, 1), offsets=offsets, labels=np.array(labels),
                        sources=np.array(["synthetic"] * len(labels)), dead=np.array(dead))
    print(len(labels), "synthetic crops")


def cmd_det_eval(argv):
    """Class-agnostic detection quality of the reference model against COCO val2017 ground truth, for several
weight variants: AP at IoU 0.5 over all proposals (score-ranked, greedy matching), and recall and precision at
a score threshold. Crowd annotations are ignored. A variant is name=path.safetensors, whose tensors replace
Mirage's float32 ones, or name=stock, Mirage's weights through the stock quantiser."""
    ap = parser("det-eval", cmd_det_eval.__doc__)
    ap.add_argument("--n", type=int, default=100)
    ap.add_argument("--first", type=int, default=0)
    ap.add_argument("--variant", action="append", default=[])
    ap.add_argument("--thr", type=float, default=0.4)
    args = ap.parse_args(argv)
    model = Mirage(mirage_weights(args))
    p32 = {k: v.clone() for k, v in model.p.items()}
    arc = Archive(mirage_archive(args))
    plan = cq_plan()
    gt = coco_truth(args)
    photos = val300(args, args.first, args.n)
    variants = [("fp32", None)] + [tuple(v.split("=", 1)) for v in args.variant]
    for name, src in variants:
        if src is None:
            model.p = {k: v.clone() for k, v in p32.items()}
        elif src == "stock":
            model.p, _ = variant(p32, plan, arc.codebook, "stock")
        else:
            model.p = dict(p32)
            model.p.update({k: v.float() for k, v in load_file(src).items()})
        model.T = model.p["embedding/embedding"]
        dets, n_gt = [], 0
        hits_thr = props_thr = matched_thr = 0
        t0 = time.time()
        for iid, path in photos:
            img = Image.open(path).convert("RGB")
            enc = model.encode(img)
            det = model.detect(enc, conf=0.05, text_conf=1.1, nms_iou=0.5)
            meta = enc["meta"]
            g = gt.get(int(iid), [])
            n_gt += len(g)
            boxes = []
            for b in det["object"]:
                bb = model.to_1000(b["box"], meta)
                boxes.append(([bb[0] / 1000 * meta["W"], bb[1] / 1000 * meta["H"], bb[2] / 1000 * meta["W"], bb[3] / 1000 * meta["H"]], b["score"]))
            used = set()
            for box, s in sorted(boxes, key=lambda x: -x[1]):
                best, bi = 0.0, -1
                for k, gb in enumerate(g):
                    if k in used:
                        continue
                    v = iou(box, gb)
                    if v > best:
                        best, bi = v, k
                tp = best >= 0.5
                if tp:
                    used.add(bi)
                dets.append((s, tp))
            sel = [bx for bx, s in boxes if s >= args.thr]
            props_thr += len(sel)
            matched_thr += sum(any(iou(bx, gb) >= 0.5 for gb in g) for bx in sel)
            hits_thr += sum(any(iou(bx, gb) >= 0.5 for bx in sel) for gb in g)
        dets.sort(key=lambda x: -x[0])
        tp = np.cumsum([d[1] for d in dets])
        fp = np.cumsum([not d[1] for d in dets])
        rec = tp / max(n_gt, 1)
        prec = tp / np.maximum(tp + fp, 1)
        ap50 = 0.0
        for t in np.linspace(0, 1, 101):
            p = prec[rec >= t]
            ap50 += (p.max() if p.size else 0.0) / 101
        print(f"{name:10s} images {len(photos)} gt {n_gt} AP50 {ap50:.4f} max_recall {rec[-1] if len(rec) else 0:.4f} "
              f"| at score>={args.thr}: recall {hits_thr / max(n_gt, 1):.4f} precision {matched_thr / max(props_thr, 1):.4f} "
              f"proposals/img {props_thr / len(photos):.1f} | {time.time() - t0:.0f}s", flush=True)

# ===================================================================================== entry point
COMMANDS = {
    "export": (cmd_export, "write an .oasis weight file for the engine"),
    "checkpoint": (cmd_checkpoint, "assemble a float32 checkpoint from Mirage's weights and replacements"),
    "run": (cmd_run, "the reference forward pass on one image"),
    "det-eval": (cmd_det_eval, "detection quality of weight variants on the reference model"),
    "train": (cmd_train, "quantisation-aware training of the text reader"),
    "eval-reader": (cmd_eval_reader, "score readers on packed word crops"),
    "eval-detected": (cmd_eval_detected, "score readers on the engine's detector crops"),
    "pack-words": (cmd_pack_words, "pack word-crop benchmarks into an archive"),
    "pack-detected": (cmd_pack_detected, "pack the crops the engine's text detector cuts"),
    "synth-val": (cmd_synth_val, "write a synthetic validation set"),
}


def main(argv=None):
    argv = sys.argv[1:] if argv is None else argv
    if not argv or argv[0] not in COMMANDS:
        print(__doc__.rstrip())
        print()
        for name, (_, what) in COMMANDS.items():
            print(f"  {name:14s} {what}")
        sys.exit(0 if not argv or argv[0] in ("-h", "--help") else 2)
    COMMANDS[argv[0]][0](argv[1:])


if __name__ == "__main__":
    main()