Nadir v2: element head, inner-page data, runtime fixes

864f427 verified 21 days ago

20.3 kB

#!/usr/bin/env python3
"""Train NadirNet on Haiku-designed and native dark-mode targets, and its element head on Haiku's image labels."""
import argparse, json, math, os, pickle, random, time
from collections import defaultdict
import numpy as np
import torch
import torch.nn.functional as F

from model import NadirNet, K_REL, FI, N_TREAT, N_IROLE, contrast_ratio
from build_ds import oklab

COLOR_IDX = dict(lab=(0, 1, 2), C=3, alpha=4, cs=(5, 6), rgb=(7, 8, 9), under=(34, 35, 36), over=(39, 40, 41))


def lab_to_rgb_np(lab):
    M1 = np.array([[1.0, 0.3963377774, 0.2158037573], [1.0, -0.1055613458, -0.0638541728], [1.0, -0.0894841775, -1.2914855480]])
    M2 = np.array([[4.0767416621, -3.3077115913, 0.2309699292], [-1.2684380046, 2.6097574011, -0.3413193965], [-0.0041960863, -0.7034186147, 1.7076147010]])
    lin = np.clip((lab @ M1.T) ** 3 @ M2.T, 0, 1)
    return np.where(lin <= 0.0031308, 12.92 * lin, 1.055 * np.power(lin, 1 / 2.4) - 0.055)


def prep(s):
    """Precompute OKLab of inputs/targets once per sample."""
    if "lab_in" in s:
        return s
    s["lab_in"] = oklab(s["rgba"][:, :3]).astype(np.float32)
    t = s["tgt"]
    ok = np.isfinite(t[:, 0])
    lab_t = np.zeros((len(t), 3), np.float32)
    if ok.any():
        lab_t[ok] = oklab(t[ok, :3]).astype(np.float32)
    s["lab_t"], s["tmask"] = lab_t, ok
    s["alpha_t"] = np.where(ok, t[:, 3], 1).astype(np.float32)
    km = np.zeros(len(t), bool)
    if len(s["keep"]):
        km[s["keep"]] = True
    s["tmask"] = s["tmask"] & ~km
    if "E_ut" not in s:
        s.update(E_f=np.zeros((0, FI), np.float32), E_thumb=np.zeros((0, 32, 32, 4), np.uint8), E_ul=np.zeros((0, 3), np.float32),
                 E_ut=np.zeros(0, np.int64), E_soft=np.zeros((0, N_TREAT), np.float32), E_role=np.zeros(0, np.int64))
    s["root"] = int(np.argmax(s["F"][:, 24])) if len(s["F"]) else 0
    return s


def augment(s, rng):
    s = dict(s)
    n = len(s["role"])
    # token dropout (simulates ad blockers, A/B variants, partial loads)
    if rng.random() < 0.6 and n > 8:
        p = rng.uniform(0, 0.25)
        root = s["F"][:, 24] > 0.5
        keepm = (rng.random(n) > p) | root
        idx = np.nonzero(keepm)[0]
        remap = -np.ones(n, np.int64); remap[idx] = np.arange(len(idx))
        for k in ("F", "role", "rgba", "w", "lab_in", "lab_t", "tmask", "alpha_t"):
            s[k] = s[k][idx]
        rel = s["rel"]
        if len(rel):
            a, b = rel[:, 1].astype(int), rel[:, 2].astype(int)
            good = (remap[a] >= 0) & (remap[b] >= 0)
            rel = rel[good].copy(); rel[:, 1] = remap[a[good]]; rel[:, 2] = remap[b[good]]
        s["rel"] = rel
        iu = s["img_under"].copy(); m = iu >= 0; iu[m] = remap[iu[m]]; s["img_under"] = iu
        eu = s["E_ut"].copy(); eu[eu >= len(remap)] = -1; m = eu >= 0; eu[m] = remap[eu[m]]; s["E_ut"] = eu
        s["root"] = int(remap[s["root"]]) if remap[s["root"]] >= 0 else 0
        n = len(idx)
    # hue rotation of the whole design (inputs and targets together)
    if rng.random() < 0.3:
        th = rng.uniform(0, 2 * math.pi)
        c, si = math.cos(th), math.sin(th)
        rot = lambda a, b: (c * a - si * b, si * a + c * b)
        F_ = s["F"].copy()
        for (ia, ib) in ((1, 2), (35, 36), (40, 41), (5, 6)):
            F_[:, ia], F_[:, ib] = rot(F_[:, ia], F_[:, ib])
        lab = s["lab_in"].copy(); lab[:, 1], lab[:, 2] = rot(lab[:, 1], lab[:, 2])
        F_[:, 7:10] = lab_to_rgb_np(lab.astype(np.float64)).astype(np.float32)
        s["F"], s["lab_in"] = F_, lab.astype(np.float32)
        lt = s["lab_t"].copy(); lt[:, 1], lt[:, 2] = rot(lt[:, 1], lt[:, 2]); s["lab_t"] = lt.astype(np.float32)
        # the surfaces under images turn with the design (the images themselves stay as they are)
        for k in ("IF", "E_f"):
            if len(s[k]):
                x = s[k].copy(); x[:, 24], x[:, 25] = rot(x[:, 24], x[:, 25]); s[k] = x
        if len(s["E_ul"]):
            u = s["E_ul"].copy(); u[:, 1], u[:, 2] = rot(u[:, 1], u[:, 2]); s["E_ul"] = u.astype(np.float32)
    return s


def collate(batch, device):
    B = len(batch)
    N = max(1, max(len(s["role"]) for s in batch))
    M = max(len(s["IF"]) for s in batch)
    T = 1 + N + M
    tok_f = np.zeros((B, N, 49), np.float32); role = np.zeros((B, N), np.int64); lab = np.zeros((B, N, 3), np.float32)
    lab_t = np.zeros((B, N, 3), np.float32); tmask = np.zeros((B, N), bool); w = np.zeros((B, N), np.float32)
    alpha_in = np.ones((B, N), np.float32); alpha_t = np.ones((B, N), np.float32)
    img_f = np.zeros((B, M, FI), np.float32); thumbs = np.zeros((B, M, 4, 32, 32), np.float32); imgt = -np.ones((B, M), np.int64)
    K = max(1, max(min(64, len(s["E_ut"])) for s in batch))
    e_f = np.zeros((B, K, FI), np.float32); e_thumb = np.zeros((B, K, 4, 32, 32), np.float32); e_ul = np.zeros((B, K, 3), np.float32)
    e_ut = np.zeros((B, K), np.int64); e_soft = np.zeros((B, K, N_TREAT), np.float32); e_role = -np.ones((B, K), np.int64); e_mask = np.zeros((B, K), bool)
    page = np.zeros((B, 8), np.float32); src = np.zeros(B, np.int64)
    rel = np.zeros((B, T, T, K_REL), np.float32); mask = np.zeros((B, T), bool)
    tb_pair = -np.ones((B, N), np.int64)
    for i, s in enumerate(batch):
        n, m = len(s["role"]), len(s["IF"])
        tok_f[i, :n] = s["F"]; role[i, :n] = s["role"]; lab[i, :n] = s["lab_in"]; lab_t[i, :n] = s["lab_t"]; tmask[i, :n] = s["tmask"]
        w[i, :n] = s["w"]; alpha_in[i, :n] = s["rgba"][:, 3]; alpha_t[i, :n] = s["alpha_t"]
        if m:
            img_f[i, :m] = s["IF"]; thumbs[i, :m] = s["thumbs"].transpose(0, 3, 1, 2) / 255.0; imgt[i, :m] = s["imgt"]
        page[i] = s["page"]; src[i] = s["src"]
        mask[i, :1 + n] = True; mask[i, 1 + N:1 + N + m] = True
        k = len(s["E_ut"])
        if k:
            sel = np.arange(k) if k <= K else np.sort(np.random.default_rng(len(s["role"]) * 7 + k).choice(k, K, replace=False))
            k = len(sel)
            e_f[i, :k] = s["E_f"][sel]; e_thumb[i, :k] = s["E_thumb"][sel].transpose(0, 3, 1, 2) / 255.0; e_ul[i, :k] = s["E_ul"][sel]
            ut = s["E_ut"][sel]
            e_ut[i, :k] = np.where((ut >= 0) & (ut < n), ut, s["root"]); e_soft[i, :k] = s["E_soft"][sel]; e_role[i, :k] = s["E_role"][sel]; e_mask[i, :k] = True
        # relations
        r = s["rel"]
        if len(r):
            typ, a, b, wt = r[:, 0].astype(int), r[:, 1].astype(int) + 1, r[:, 2].astype(int) + 1, r[:, 3]
            R = np.zeros((4, n + 1, n + 1), np.float32)
            np.maximum.at(R, (typ, a, b), wt)
            R /= R.max(axis=2, keepdims=True) + 1e-6
            rel[i, :n + 1, :n + 1, 0:4] = R.transpose(1, 2, 0)
            rel[i, :n + 1, :n + 1, 4:8] = R.transpose(2, 1, 0)
            # dominant surface for each text token (contrast loss)
            tbm = typ == 0
            best = {}
            for aa, bb, ww in zip(a[tbm] - 1, b[tbm] - 1, wt[tbm]):
                if ww > best.get(aa, (-1, -1))[0]:
                    best[aa] = (ww, bb)
            for aa, (ww, bb) in best.items():
                tb_pair[i, aa] = bb
        for j in range(m):
            u = s["img_under"][j]
            if u >= 0:
                rel[i, 1 + N + j, 1 + u, 8] = 1.0
                rel[i, 1 + u, 1 + N + j, 9] = 1.0
        # same colour, different role
        key = s["rgba"][:, :3] @ np.array([65536, 256, 1], np.float32)
        same = (key[:, None] == key[None, :]) & ~np.eye(n, dtype=bool)
        rel[i, 1:1 + n, 1:1 + n, 10] = same
    out = dict(tok_f=tok_f, role=role, lab=lab, lab_t=lab_t, tmask=tmask, w=w, alpha_in=alpha_in, alpha_t=alpha_t,
               img_f=img_f, thumbs=thumbs, imgt=imgt, page=page, src=src, rel=rel, mask=mask, tb_pair=tb_pair,
               e_f=e_f, e_thumb=e_thumb, e_ul=e_ul, e_ut=e_ut, e_soft=e_soft, e_role=e_role, e_mask=e_mask)
    if device is None:
        return out
    return {k: torch.from_numpy(v).to(device, non_blocking=True) for k, v in out.items()}


class Batches(torch.utils.data.Dataset):
    """Each item is a whole (augmented) batch, assembled in a worker process."""
    def __init__(self, data, order, bs, aug, seed):
        self.data, self.order, self.bs, self.aug, self.seed = data, order, bs, aug, seed

    def __len__(self):
        return math.ceil(len(self.order) / self.bs)

    def __getitem__(self, k):
        rng = np.random.default_rng((self.seed, k))
        items = [self.data[i] for i in self.order[k * self.bs:(k + 1) * self.bs]]
        if self.aug:
            items = [augment(s, rng) for s in items]
        return collate(items, None)


def to_dev(b, device):
    return {k: (v if torch.is_tensor(v) else torch.from_numpy(v)).to(device, non_blocking=True) for k, v in b.items()}


def huber(x, d):
    ax = x.abs()
    return torch.where(ax < d, 0.5 * x * x / d, ax - 0.5 * d)


def element_logits(model, b, lab, ctx):
    """The element head reads the dark surface the palette head designed under each element (as the runtime does)."""
    under_dark = torch.gather(lab.detach(), 1, b["e_ut"][..., None].expand(-1, -1, 3))
    t, r = model.elements(b["e_f"], b["e_thumb"], b["e_ul"], under_dark, ctx)
    return t.float(), r.float()


def losses(model, b, img_w=None):
    lab, alpha_logit, ctx = model(b["tok_f"], b["role"], b["lab"], b["img_f"], b["thumbs"], b["page"], b["src"], b["rel"], b["mask"])
    lab = lab.float(); alpha_logit = alpha_logit.float()
    m = b["tmask"].float() * b["w"]
    d = lab - b["lab_t"]
    per = 2.0 * huber(d[..., 0], 0.02) + huber(d[..., 1], 0.01) + huber(d[..., 2], 0.01)
    lc = (per * m).sum() / m.sum().clamp_min(1)
    # alpha (only where translucency is involved)
    am = m * ((b["alpha_in"] < 0.999) | (b["alpha_t"] < 0.999)).float()
    a_in = b["alpha_in"].clamp(1e-3, 1 - 1e-3)
    alpha = torch.sigmoid(alpha_logit + torch.log(a_in / (1 - a_in)))
    la = (huber(alpha - b["alpha_t"], 0.05) * am).sum() / am.sum().clamp_min(1)
    # readability: text must not lose contrast against its own surface
    pair = b["tb_pair"]
    pm = (pair >= 0) & b["tmask"] & (b["role"] == 1) & (b["alpha_in"] > 0.99)
    if pm.any():
        idx = pair.clamp_min(0)
        bg_pred = torch.gather(lab, 1, idx[..., None].expand(-1, -1, 3))
        bg_t = torch.gather(b["lab_t"], 1, idx[..., None].expand(-1, -1, 3))
        bgm = torch.gather(b["tmask"], 1, idx)
        pm = pm & bgm
        c_pred = contrast_ratio(lab, bg_pred)
        c_t = contrast_ratio(b["lab_t"], bg_t).clamp(max=4.5)
        lcon = (F.relu(c_t - c_pred) * pm.float()).sum() / pm.float().sum().clamp_min(1)
    else:
        lcon = lab.new_zeros(())
    # elements: soft targets (best rendering 0.75, other acceptable ones share 0.25) and the image's role
    em = b["e_mask"]
    if em.any():
        tl, rl = element_logits(model, b, lab, ctx)
        soft = b["e_soft"][em]
        lp = F.log_softmax(tl[em], -1)
        le = -(soft * lp * (img_w if img_w is not None else 1.0)).sum(-1).mean()
        rm = em & (b["e_role"] >= 0)
        lr = F.cross_entropy(rl[rm], b["e_role"][rm], label_smoothing=0.05) if rm.any() else lab.new_zeros(())
        with torch.no_grad():
            pred = tl[em].argmax(-1)
            e_acc = (pred == soft.argmax(-1)).float().mean().item()
            e_ok = (soft.gather(1, pred[:, None])[:, 0] > 0).float().mean().item()
            r_acc = (rl[rm].argmax(-1) == b["e_role"][rm]).float().mean().item() if rm.any() else 0.0
        n_el, n_role = int(em.sum()), int(rm.sum())
    else:
        le = lr = lab.new_zeros(()); e_acc = e_ok = r_acc = 0.0; n_el = n_role = 0
    with torch.no_grad():
        de = (d.norm(dim=-1) * m).sum() / m.sum().clamp_min(1)
        dl = (d[..., 0].abs() * m).sum() / m.sum().clamp_min(1)
    return lc, la, lcon, le, lr, dict(dE=de.item(), dL=dl.item(), e_acc=e_acc, e_ok=e_ok, r_acc=r_acc, n_el=n_el, n_role=n_role)


def evaluate(model, val, device, bs, img_w):
    model.eval()
    agg = defaultdict(float); cnt = defaultdict(float)
    with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
        for src in (0, 1):
            vs = [s for s in val if s["src"] == src]
            for k in range(0, len(vs), bs):
                b = collate(vs[k:k + bs], device)
                lc, la, lcon, le, lr, st = losses(model, b, img_w)
                n = len(vs[k:k + bs])
                agg[f"s{src}_dE"] += st["dE"] * n; agg[f"s{src}_dL"] += st["dL"] * n; agg[f"s{src}_con"] += lcon.item() * n
                cnt[f"s{src}"] += n
                if st["n_el"]:
                    for key in ("e_acc", "e_ok"):
                        agg[key] += st[key] * st["n_el"]
                    agg["e_nll"] += le.item() * st["n_el"]; cnt["e"] += st["n_el"]
                if st["n_role"]:
                    agg["r_acc"] += st["r_acc"] * st["n_role"]; cnt["r"] += st["n_role"]
    out = {}
    for k, v in agg.items():
        if k.startswith("e_"):
            out[k] = v / max(1, cnt["e"])
        elif k.startswith("r_"):
            out[k] = v / max(1, cnt["r"])
        else:
            out[k] = v / max(1, cnt[k.split("_")[0]])
    model.train()
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--d", type=int, default=384)
    ap.add_argument("--layers", type=int, default=8)
    ap.add_argument("--heads", type=int, default=6)
    ap.add_argument("--epochs", type=int, default=40)
    ap.add_argument("--bs", type=int, default=48)
    ap.add_argument("--lr", type=float, default=4e-4)
    ap.add_argument("--wd", type=float, default=0.05)
    ap.add_argument("--drop", type=float, default=0.1)
    ap.add_argument("--native-weight", type=float, default=0.5)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--workers", type=int, default=8)
    ap.add_argument("--elem-weight", type=float, default=0.1)
    ap.add_argument("--role-weight", type=float, default=0.03)
    ap.add_argument("--elem-balance", type=float, default=0.3, help="element class weights = (mass share)^-balance")
    ap.add_argument("--init", default="", help="warm start from a checkpoint (matching tensors; widened inputs are zero-padded)")
    args = ap.parse_args()
    torch.manual_seed(args.seed); rng = np.random.default_rng(args.seed)
    os.makedirs(args.out, exist_ok=True)
    data = pickle.load(open(args.data, "rb"))
    train, val = [prep(s) for s in data["train"]], [prep(s) for s in data["val"]]
    print(f"train {len(train)} val {len(val)} (haiku {sum(s['src'] == 0 for s in train)}, native {sum(s['src'] == 1 for s in train)})", flush=True)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = NadirNet(d=args.d, layers=args.layers, heads=args.heads, drop=args.drop).to(device)
    nparam = sum(p.numel() for p in model.parameters())
    print(f"params {nparam / 1e6:.2f}M", flush=True)
    if args.init:
        src_sd = torch.load(args.init, map_location="cpu")["model"]
        sd = model.state_dict()
        loaded = 0
        for k, v in src_sd.items():
            if k not in sd:
                continue
            if sd[k].shape == v.shape:
                sd[k] = v; loaded += 1
            elif sd[k].dim() == 2 and sd[k].shape[0] == v.shape[0] and sd[k].shape[1] > v.shape[1]:
                w = torch.zeros_like(sd[k]); w[:, :v.shape[1]] = v; sd[k] = w; loaded += 1
        model.load_state_dict(sd)
        print(f"warm start from {args.init}: {loaded} tensors", flush=True)
    img_w = None
    el = np.zeros(N_TREAT)
    for s in train:
        el += s["E_soft"].sum(0) if len(s["E_soft"]) else 0
    print("element label mass", el.round().astype(int).tolist(), "pages with elements", sum(len(s["E_ut"]) > 0 for s in train), flush=True)
    if el.sum() > 0:
        cw = (np.maximum(el, 1) / el.sum()) ** (-args.elem_balance)
        img_w = torch.tensor(cw / cw.mean(), dtype=torch.float32, device=device)
        print("element class weights", [round(float(x), 3) for x in img_w], flush=True)
    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.wd, betas=(0.9, 0.98))
    steps_per = math.ceil(len(train) / args.bs)
    total = steps_per * args.epochs
    warm = min(800, total // 20)
    sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda st: min(1, (st + 1) / warm) * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * min(1, st / total)))))
    # sample weights: native pages count less (different designers, noisier pairing)
    sw = np.array([1.0 if s["src"] == 0 else args.native_weight for s in train]); sw /= sw.sum()
    best = 1e9; step = 0; log = []
    t0 = time.time()
    for ep in range(args.epochs):
        order = rng.choice(len(train), size=len(train), p=sw)
        run = defaultdict(float); nb = 0
        loader = torch.utils.data.DataLoader(Batches(train, order, args.bs, True, args.seed * 1000 + ep), batch_size=None, shuffle=False,
                                             num_workers=args.workers, persistent_workers=False, prefetch_factor=4 if args.workers else None)
        for bnp in loader:
            b = to_dev(bnp, device)
            with torch.autocast("cuda", dtype=torch.bfloat16):
                lc, la, lcon, le, lr, st = losses(model, b, img_w)
            loss = lc + 0.5 * la + 0.02 * lcon + args.elem_weight * le + args.role_weight * lr
            opt.zero_grad(set_to_none=True)
            if not torch.isfinite(loss):  # one bad page must never poison the weights
                run["skipped"] += 1
                continue
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step(); sched.step(); step += 1
            run["loss"] += loss.item(); run["dE"] += st["dE"]; run["con"] += lcon.item(); run["elem"] += le.item(); run["e_ok"] += st["e_ok"]; nb += 1
        ev = evaluate(model, val, device, 64, img_w)
        rec = {"epoch": ep, "step": step, "secs": round(time.time() - t0), **{k: round(v / nb, 5) for k, v in run.items()}, **{k: round(v, 5) for k, v in ev.items()}}
        log.append(rec)
        print(json.dumps(rec), flush=True)
        score = ev.get("s0_dE", 9) + 0.3 * ev.get("s1_dE", 0) + 0.02 * ev.get("s0_con", 0) + 0.05 * (1 - ev.get("e_ok", 0))
        if score < best:
            best = score
            torch.save({"model": model.state_dict(), "cfg": model.cfg, "epoch": ep, "eval": ev}, os.path.join(args.out, "best.pt"))
    torch.save({"model": model.state_dict(), "cfg": model.cfg, "epoch": args.epochs - 1}, os.path.join(args.out, "last.pt"))
    # confidence calibration of the element head: one temperature fitted on held-out sites
    ck = torch.load(os.path.join(args.out, "best.pt"), map_location=device)
    model.load_state_dict(ck["model"]); model.eval()
    logits, softs = [], []
    with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
        vs = [s for s in val if len(s["E_ut"])]
        for k in range(0, len(vs), 64):
            b = collate(vs[k:k + 64], device)
            lab, _, ctx = model(b["tok_f"], b["role"], b["lab"], b["img_f"], b["thumbs"], b["page"], b["src"], b["rel"], b["mask"])
            tl, _ = element_logits(model, b, lab.float(), ctx)
            logits.append(tl[b["e_mask"]].float()); softs.append(b["e_soft"][b["e_mask"]])
    if logits:
        lg, sf = torch.cat(logits), torch.cat(softs)
        hard = sf.argmax(-1)
        best_t, best_nll = 1.0, 1e9
        for tt in np.arange(0.5, 3.01, 0.05):
            nll = F.cross_entropy(lg / tt, hard).item()
            if nll < best_nll:
                best_t, best_nll = float(tt), nll
        p = F.softmax(lg / best_t, -1); conf, pred = p.max(-1); corr = (pred == hard).float()
        bins = torch.clamp((conf * 10).long(), max=9)
        ece = sum((bins == i).float().mean() * (conf[bins == i].mean() - corr[bins == i].mean()).abs() for i in range(10) if (bins == i).any()).item()
        ck["temperature"] = best_t
        ck["calibration"] = {"temperature": best_t, "nll": best_nll, "ece": ece, "n": int(len(hard))}
        torch.save(ck, os.path.join(args.out, "best.pt"))
        print("element head calibration", ck["calibration"], flush=True)
    json.dump({"args": vars(args), "params": nparam, "log": log, "best_score": best}, open(os.path.join(args.out, "train_log.json"), "w"), indent=1)


if __name__ == "__main__":
    main()

	#!/usr/bin/env python3
	"""Train NadirNet on Haiku-designed and native dark-mode targets, and its element head on Haiku's image labels."""
	import argparse, json, math, os, pickle, random, time
	from collections import defaultdict
	import numpy as np
	import torch
	import torch.nn.functional as F

	from model import NadirNet, K_REL, FI, N_TREAT, N_IROLE, contrast_ratio
	from build_ds import oklab

	COLOR_IDX = dict(lab=(0, 1, 2), C=3, alpha=4, cs=(5, 6), rgb=(7, 8, 9), under=(34, 35, 36), over=(39, 40, 41))


	def lab_to_rgb_np(lab):
	M1 = np.array([[1.0, 0.3963377774, 0.2158037573], [1.0, -0.1055613458, -0.0638541728], [1.0, -0.0894841775, -1.2914855480]])
	M2 = np.array([[4.0767416621, -3.3077115913, 0.2309699292], [-1.2684380046, 2.6097574011, -0.3413193965], [-0.0041960863, -0.7034186147, 1.7076147010]])
	lin = np.clip((lab @ M1.T) ** 3 @ M2.T, 0, 1)
	return np.where(lin <= 0.0031308, 12.92 * lin, 1.055 * np.power(lin, 1 / 2.4) - 0.055)


	def prep(s):
	"""Precompute OKLab of inputs/targets once per sample."""
	if "lab_in" in s:
	return s
	s["lab_in"] = oklab(s["rgba"][:, :3]).astype(np.float32)
	t = s["tgt"]
	ok = np.isfinite(t[:, 0])
	lab_t = np.zeros((len(t), 3), np.float32)
	if ok.any():
	lab_t[ok] = oklab(t[ok, :3]).astype(np.float32)
	s["lab_t"], s["tmask"] = lab_t, ok
	s["alpha_t"] = np.where(ok, t[:, 3], 1).astype(np.float32)
	km = np.zeros(len(t), bool)
	if len(s["keep"]):
	km[s["keep"]] = True
	s["tmask"] = s["tmask"] & ~km
	if "E_ut" not in s:
	s.update(E_f=np.zeros((0, FI), np.float32), E_thumb=np.zeros((0, 32, 32, 4), np.uint8), E_ul=np.zeros((0, 3), np.float32),
	E_ut=np.zeros(0, np.int64), E_soft=np.zeros((0, N_TREAT), np.float32), E_role=np.zeros(0, np.int64))
	s["root"] = int(np.argmax(s["F"][:, 24])) if len(s["F"]) else 0
	return s


	def augment(s, rng):
	s = dict(s)
	n = len(s["role"])
	# token dropout (simulates ad blockers, A/B variants, partial loads)
	if rng.random() < 0.6 and n > 8:
	p = rng.uniform(0, 0.25)
	root = s["F"][:, 24] > 0.5
	keepm = (rng.random(n) > p) \| root
	idx = np.nonzero(keepm)[0]
	remap = -np.ones(n, np.int64); remap[idx] = np.arange(len(idx))
	for k in ("F", "role", "rgba", "w", "lab_in", "lab_t", "tmask", "alpha_t"):
	s[k] = s[k][idx]
	rel = s["rel"]
	if len(rel):
	a, b = rel[:, 1].astype(int), rel[:, 2].astype(int)
	good = (remap[a] >= 0) & (remap[b] >= 0)
	rel = rel[good].copy(); rel[:, 1] = remap[a[good]]; rel[:, 2] = remap[b[good]]
	s["rel"] = rel
	iu = s["img_under"].copy(); m = iu >= 0; iu[m] = remap[iu[m]]; s["img_under"] = iu
	eu = s["E_ut"].copy(); eu[eu >= len(remap)] = -1; m = eu >= 0; eu[m] = remap[eu[m]]; s["E_ut"] = eu
	s["root"] = int(remap[s["root"]]) if remap[s["root"]] >= 0 else 0
	n = len(idx)
	# hue rotation of the whole design (inputs and targets together)
	if rng.random() < 0.3:
	th = rng.uniform(0, 2 * math.pi)
	c, si = math.cos(th), math.sin(th)
	rot = lambda a, b: (c * a - si * b, si * a + c * b)
	F_ = s["F"].copy()
	for (ia, ib) in ((1, 2), (35, 36), (40, 41), (5, 6)):
	F_[:, ia], F_[:, ib] = rot(F_[:, ia], F_[:, ib])
	lab = s["lab_in"].copy(); lab[:, 1], lab[:, 2] = rot(lab[:, 1], lab[:, 2])
	F_[:, 7:10] = lab_to_rgb_np(lab.astype(np.float64)).astype(np.float32)
	s["F"], s["lab_in"] = F_, lab.astype(np.float32)
	lt = s["lab_t"].copy(); lt[:, 1], lt[:, 2] = rot(lt[:, 1], lt[:, 2]); s["lab_t"] = lt.astype(np.float32)
	# the surfaces under images turn with the design (the images themselves stay as they are)
	for k in ("IF", "E_f"):
	if len(s[k]):
	x = s[k].copy(); x[:, 24], x[:, 25] = rot(x[:, 24], x[:, 25]); s[k] = x
	if len(s["E_ul"]):
	u = s["E_ul"].copy(); u[:, 1], u[:, 2] = rot(u[:, 1], u[:, 2]); s["E_ul"] = u.astype(np.float32)
	return s


	def collate(batch, device):
	B = len(batch)
	N = max(1, max(len(s["role"]) for s in batch))
	M = max(len(s["IF"]) for s in batch)
	T = 1 + N + M
	tok_f = np.zeros((B, N, 49), np.float32); role = np.zeros((B, N), np.int64); lab = np.zeros((B, N, 3), np.float32)
	lab_t = np.zeros((B, N, 3), np.float32); tmask = np.zeros((B, N), bool); w = np.zeros((B, N), np.float32)
	alpha_in = np.ones((B, N), np.float32); alpha_t = np.ones((B, N), np.float32)
	img_f = np.zeros((B, M, FI), np.float32); thumbs = np.zeros((B, M, 4, 32, 32), np.float32); imgt = -np.ones((B, M), np.int64)
	K = max(1, max(min(64, len(s["E_ut"])) for s in batch))
	e_f = np.zeros((B, K, FI), np.float32); e_thumb = np.zeros((B, K, 4, 32, 32), np.float32); e_ul = np.zeros((B, K, 3), np.float32)
	e_ut = np.zeros((B, K), np.int64); e_soft = np.zeros((B, K, N_TREAT), np.float32); e_role = -np.ones((B, K), np.int64); e_mask = np.zeros((B, K), bool)
	page = np.zeros((B, 8), np.float32); src = np.zeros(B, np.int64)
	rel = np.zeros((B, T, T, K_REL), np.float32); mask = np.zeros((B, T), bool)
	tb_pair = -np.ones((B, N), np.int64)
	for i, s in enumerate(batch):
	n, m = len(s["role"]), len(s["IF"])
	tok_f[i, :n] = s["F"]; role[i, :n] = s["role"]; lab[i, :n] = s["lab_in"]; lab_t[i, :n] = s["lab_t"]; tmask[i, :n] = s["tmask"]
	w[i, :n] = s["w"]; alpha_in[i, :n] = s["rgba"][:, 3]; alpha_t[i, :n] = s["alpha_t"]
	if m:
	img_f[i, :m] = s["IF"]; thumbs[i, :m] = s["thumbs"].transpose(0, 3, 1, 2) / 255.0; imgt[i, :m] = s["imgt"]
	page[i] = s["page"]; src[i] = s["src"]
	mask[i, :1 + n] = True; mask[i, 1 + N:1 + N + m] = True
	k = len(s["E_ut"])
	if k:
	sel = np.arange(k) if k <= K else np.sort(np.random.default_rng(len(s["role"]) * 7 + k).choice(k, K, replace=False))
	k = len(sel)
	e_f[i, :k] = s["E_f"][sel]; e_thumb[i, :k] = s["E_thumb"][sel].transpose(0, 3, 1, 2) / 255.0; e_ul[i, :k] = s["E_ul"][sel]
	ut = s["E_ut"][sel]
	e_ut[i, :k] = np.where((ut >= 0) & (ut < n), ut, s["root"]); e_soft[i, :k] = s["E_soft"][sel]; e_role[i, :k] = s["E_role"][sel]; e_mask[i, :k] = True
	# relations
	r = s["rel"]
	if len(r):
	typ, a, b, wt = r[:, 0].astype(int), r[:, 1].astype(int) + 1, r[:, 2].astype(int) + 1, r[:, 3]
	R = np.zeros((4, n + 1, n + 1), np.float32)
	np.maximum.at(R, (typ, a, b), wt)
	R /= R.max(axis=2, keepdims=True) + 1e-6
	rel[i, :n + 1, :n + 1, 0:4] = R.transpose(1, 2, 0)
	rel[i, :n + 1, :n + 1, 4:8] = R.transpose(2, 1, 0)
	# dominant surface for each text token (contrast loss)
	tbm = typ == 0
	best = {}
	for aa, bb, ww in zip(a[tbm] - 1, b[tbm] - 1, wt[tbm]):
	if ww > best.get(aa, (-1, -1))[0]:
	best[aa] = (ww, bb)
	for aa, (ww, bb) in best.items():
	tb_pair[i, aa] = bb
	for j in range(m):
	u = s["img_under"][j]
	if u >= 0:
	rel[i, 1 + N + j, 1 + u, 8] = 1.0
	rel[i, 1 + u, 1 + N + j, 9] = 1.0
	# same colour, different role
	key = s["rgba"][:, :3] @ np.array([65536, 256, 1], np.float32)
	same = (key[:, None] == key[None, :]) & ~np.eye(n, dtype=bool)
	rel[i, 1:1 + n, 1:1 + n, 10] = same
	out = dict(tok_f=tok_f, role=role, lab=lab, lab_t=lab_t, tmask=tmask, w=w, alpha_in=alpha_in, alpha_t=alpha_t,
	img_f=img_f, thumbs=thumbs, imgt=imgt, page=page, src=src, rel=rel, mask=mask, tb_pair=tb_pair,
	e_f=e_f, e_thumb=e_thumb, e_ul=e_ul, e_ut=e_ut, e_soft=e_soft, e_role=e_role, e_mask=e_mask)
	if device is None:
	return out
	return {k: torch.from_numpy(v).to(device, non_blocking=True) for k, v in out.items()}


	class Batches(torch.utils.data.Dataset):
	"""Each item is a whole (augmented) batch, assembled in a worker process."""
	def __init__(self, data, order, bs, aug, seed):
	self.data, self.order, self.bs, self.aug, self.seed = data, order, bs, aug, seed

	def __len__(self):
	return math.ceil(len(self.order) / self.bs)

	def __getitem__(self, k):
	rng = np.random.default_rng((self.seed, k))
	items = [self.data[i] for i in self.order[k * self.bs:(k + 1) * self.bs]]
	if self.aug:
	items = [augment(s, rng) for s in items]
	return collate(items, None)


	def to_dev(b, device):
	return {k: (v if torch.is_tensor(v) else torch.from_numpy(v)).to(device, non_blocking=True) for k, v in b.items()}


	def huber(x, d):
	ax = x.abs()
	return torch.where(ax < d, 0.5 * x * x / d, ax - 0.5 * d)


	def element_logits(model, b, lab, ctx):
	"""The element head reads the dark surface the palette head designed under each element (as the runtime does)."""
	under_dark = torch.gather(lab.detach(), 1, b["e_ut"][..., None].expand(-1, -1, 3))
	t, r = model.elements(b["e_f"], b["e_thumb"], b["e_ul"], under_dark, ctx)
	return t.float(), r.float()


	def losses(model, b, img_w=None):
	lab, alpha_logit, ctx = model(b["tok_f"], b["role"], b["lab"], b["img_f"], b["thumbs"], b["page"], b["src"], b["rel"], b["mask"])
	lab = lab.float(); alpha_logit = alpha_logit.float()
	m = b["tmask"].float() * b["w"]
	d = lab - b["lab_t"]
	per = 2.0 * huber(d[..., 0], 0.02) + huber(d[..., 1], 0.01) + huber(d[..., 2], 0.01)
	lc = (per * m).sum() / m.sum().clamp_min(1)
	# alpha (only where translucency is involved)
	am = m * ((b["alpha_in"] < 0.999) \| (b["alpha_t"] < 0.999)).float()
	a_in = b["alpha_in"].clamp(1e-3, 1 - 1e-3)
	alpha = torch.sigmoid(alpha_logit + torch.log(a_in / (1 - a_in)))
	la = (huber(alpha - b["alpha_t"], 0.05) * am).sum() / am.sum().clamp_min(1)
	# readability: text must not lose contrast against its own surface
	pair = b["tb_pair"]
	pm = (pair >= 0) & b["tmask"] & (b["role"] == 1) & (b["alpha_in"] > 0.99)
	if pm.any():
	idx = pair.clamp_min(0)
	bg_pred = torch.gather(lab, 1, idx[..., None].expand(-1, -1, 3))
	bg_t = torch.gather(b["lab_t"], 1, idx[..., None].expand(-1, -1, 3))
	bgm = torch.gather(b["tmask"], 1, idx)
	pm = pm & bgm
	c_pred = contrast_ratio(lab, bg_pred)
	c_t = contrast_ratio(b["lab_t"], bg_t).clamp(max=4.5)
	lcon = (F.relu(c_t - c_pred) * pm.float()).sum() / pm.float().sum().clamp_min(1)
	else:
	lcon = lab.new_zeros(())
	# elements: soft targets (best rendering 0.75, other acceptable ones share 0.25) and the image's role
	em = b["e_mask"]
	if em.any():
	tl, rl = element_logits(model, b, lab, ctx)
	soft = b["e_soft"][em]
	lp = F.log_softmax(tl[em], -1)
	le = -(soft * lp * (img_w if img_w is not None else 1.0)).sum(-1).mean()
	rm = em & (b["e_role"] >= 0)
	lr = F.cross_entropy(rl[rm], b["e_role"][rm], label_smoothing=0.05) if rm.any() else lab.new_zeros(())
	with torch.no_grad():
	pred = tl[em].argmax(-1)
	e_acc = (pred == soft.argmax(-1)).float().mean().item()
	e_ok = (soft.gather(1, pred[:, None])[:, 0] > 0).float().mean().item()
	r_acc = (rl[rm].argmax(-1) == b["e_role"][rm]).float().mean().item() if rm.any() else 0.0
	n_el, n_role = int(em.sum()), int(rm.sum())
	else:
	le = lr = lab.new_zeros(()); e_acc = e_ok = r_acc = 0.0; n_el = n_role = 0
	with torch.no_grad():
	de = (d.norm(dim=-1) * m).sum() / m.sum().clamp_min(1)
	dl = (d[..., 0].abs() * m).sum() / m.sum().clamp_min(1)
	return lc, la, lcon, le, lr, dict(dE=de.item(), dL=dl.item(), e_acc=e_acc, e_ok=e_ok, r_acc=r_acc, n_el=n_el, n_role=n_role)


	def evaluate(model, val, device, bs, img_w):
	model.eval()
	agg = defaultdict(float); cnt = defaultdict(float)
	with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
	for src in (0, 1):
	vs = [s for s in val if s["src"] == src]
	for k in range(0, len(vs), bs):
	b = collate(vs[k:k + bs], device)
	lc, la, lcon, le, lr, st = losses(model, b, img_w)
	n = len(vs[k:k + bs])
	agg[f"s{src}_dE"] += st["dE"] * n; agg[f"s{src}_dL"] += st["dL"] * n; agg[f"s{src}_con"] += lcon.item() * n
	cnt[f"s{src}"] += n
	if st["n_el"]:
	for key in ("e_acc", "e_ok"):
	agg[key] += st[key] * st["n_el"]
	agg["e_nll"] += le.item() * st["n_el"]; cnt["e"] += st["n_el"]
	if st["n_role"]:
	agg["r_acc"] += st["r_acc"] * st["n_role"]; cnt["r"] += st["n_role"]
	out = {}
	for k, v in agg.items():
	if k.startswith("e_"):
	out[k] = v / max(1, cnt["e"])
	elif k.startswith("r_"):
	out[k] = v / max(1, cnt["r"])
	else:
	out[k] = v / max(1, cnt[k.split("_")[0]])
	model.train()
	return out


	def main():
	ap = argparse.ArgumentParser()
	ap.add_argument("--data", required=True)
	ap.add_argument("--out", required=True)
	ap.add_argument("--d", type=int, default=384)
	ap.add_argument("--layers", type=int, default=8)
	ap.add_argument("--heads", type=int, default=6)
	ap.add_argument("--epochs", type=int, default=40)
	ap.add_argument("--bs", type=int, default=48)
	ap.add_argument("--lr", type=float, default=4e-4)
	ap.add_argument("--wd", type=float, default=0.05)
	ap.add_argument("--drop", type=float, default=0.1)
	ap.add_argument("--native-weight", type=float, default=0.5)
	ap.add_argument("--seed", type=int, default=0)
	ap.add_argument("--workers", type=int, default=8)
	ap.add_argument("--elem-weight", type=float, default=0.1)
	ap.add_argument("--role-weight", type=float, default=0.03)
	ap.add_argument("--elem-balance", type=float, default=0.3, help="element class weights = (mass share)^-balance")
	ap.add_argument("--init", default="", help="warm start from a checkpoint (matching tensors; widened inputs are zero-padded)")
	args = ap.parse_args()
	torch.manual_seed(args.seed); rng = np.random.default_rng(args.seed)
	os.makedirs(args.out, exist_ok=True)
	data = pickle.load(open(args.data, "rb"))
	train, val = [prep(s) for s in data["train"]], [prep(s) for s in data["val"]]
	print(f"train {len(train)} val {len(val)} (haiku {sum(s['src'] == 0 for s in train)}, native {sum(s['src'] == 1 for s in train)})", flush=True)
	device = "cuda" if torch.cuda.is_available() else "cpu"
	model = NadirNet(d=args.d, layers=args.layers, heads=args.heads, drop=args.drop).to(device)
	nparam = sum(p.numel() for p in model.parameters())
	print(f"params {nparam / 1e6:.2f}M", flush=True)
	if args.init:
	src_sd = torch.load(args.init, map_location="cpu")["model"]
	sd = model.state_dict()
	loaded = 0
	for k, v in src_sd.items():
	if k not in sd:
	continue
	if sd[k].shape == v.shape:
	sd[k] = v; loaded += 1
	elif sd[k].dim() == 2 and sd[k].shape[0] == v.shape[0] and sd[k].shape[1] > v.shape[1]:
	w = torch.zeros_like(sd[k]); w[:, :v.shape[1]] = v; sd[k] = w; loaded += 1
	model.load_state_dict(sd)
	print(f"warm start from {args.init}: {loaded} tensors", flush=True)
	img_w = None
	el = np.zeros(N_TREAT)
	for s in train:
	el += s["E_soft"].sum(0) if len(s["E_soft"]) else 0
	print("element label mass", el.round().astype(int).tolist(), "pages with elements", sum(len(s["E_ut"]) > 0 for s in train), flush=True)
	if el.sum() > 0:
	cw = (np.maximum(el, 1) / el.sum()) ** (-args.elem_balance)
	img_w = torch.tensor(cw / cw.mean(), dtype=torch.float32, device=device)
	print("element class weights", [round(float(x), 3) for x in img_w], flush=True)
	opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.wd, betas=(0.9, 0.98))
	steps_per = math.ceil(len(train) / args.bs)
	total = steps_per * args.epochs
	warm = min(800, total // 20)
	sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda st: min(1, (st + 1) / warm) * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * min(1, st / total)))))
	# sample weights: native pages count less (different designers, noisier pairing)
	sw = np.array([1.0 if s["src"] == 0 else args.native_weight for s in train]); sw /= sw.sum()
	best = 1e9; step = 0; log = []
	t0 = time.time()
	for ep in range(args.epochs):
	order = rng.choice(len(train), size=len(train), p=sw)
	run = defaultdict(float); nb = 0
	loader = torch.utils.data.DataLoader(Batches(train, order, args.bs, True, args.seed * 1000 + ep), batch_size=None, shuffle=False,
	num_workers=args.workers, persistent_workers=False, prefetch_factor=4 if args.workers else None)
	for bnp in loader:
	b = to_dev(bnp, device)
	with torch.autocast("cuda", dtype=torch.bfloat16):
	lc, la, lcon, le, lr, st = losses(model, b, img_w)
	loss = lc + 0.5 * la + 0.02 * lcon + args.elem_weight * le + args.role_weight * lr
	opt.zero_grad(set_to_none=True)
	if not torch.isfinite(loss): # one bad page must never poison the weights
	run["skipped"] += 1
	continue
	loss.backward()
	torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
	opt.step(); sched.step(); step += 1
	run["loss"] += loss.item(); run["dE"] += st["dE"]; run["con"] += lcon.item(); run["elem"] += le.item(); run["e_ok"] += st["e_ok"]; nb += 1
	ev = evaluate(model, val, device, 64, img_w)
	rec = {"epoch": ep, "step": step, "secs": round(time.time() - t0), {k: round(v / nb, 5) for k, v in run.items()}, {k: round(v, 5) for k, v in ev.items()}}
	log.append(rec)
	print(json.dumps(rec), flush=True)
	score = ev.get("s0_dE", 9) + 0.3 * ev.get("s1_dE", 0) + 0.02 * ev.get("s0_con", 0) + 0.05 * (1 - ev.get("e_ok", 0))
	if score < best:
	best = score
	torch.save({"model": model.state_dict(), "cfg": model.cfg, "epoch": ep, "eval": ev}, os.path.join(args.out, "best.pt"))
	torch.save({"model": model.state_dict(), "cfg": model.cfg, "epoch": args.epochs - 1}, os.path.join(args.out, "last.pt"))
	# confidence calibration of the element head: one temperature fitted on held-out sites
	ck = torch.load(os.path.join(args.out, "best.pt"), map_location=device)
	model.load_state_dict(ck["model"]); model.eval()
	logits, softs = [], []
	with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
	vs = [s for s in val if len(s["E_ut"])]
	for k in range(0, len(vs), 64):
	b = collate(vs[k:k + 64], device)
	lab, _, ctx = model(b["tok_f"], b["role"], b["lab"], b["img_f"], b["thumbs"], b["page"], b["src"], b["rel"], b["mask"])
	tl, _ = element_logits(model, b, lab.float(), ctx)
	logits.append(tl[b["e_mask"]].float()); softs.append(b["e_soft"][b["e_mask"]])
	if logits:
	lg, sf = torch.cat(logits), torch.cat(softs)
	hard = sf.argmax(-1)
	best_t, best_nll = 1.0, 1e9
	for tt in np.arange(0.5, 3.01, 0.05):
	nll = F.cross_entropy(lg / tt, hard).item()
	if nll < best_nll:
	best_t, best_nll = float(tt), nll
	p = F.softmax(lg / best_t, -1); conf, pred = p.max(-1); corr = (pred == hard).float()
	bins = torch.clamp((conf * 10).long(), max=9)
	ece = sum((bins == i).float().mean() * (conf[bins == i].mean() - corr[bins == i].mean()).abs() for i in range(10) if (bins == i).any()).item()
	ck["temperature"] = best_t
	ck["calibration"] = {"temperature": best_t, "nll": best_nll, "ece": ece, "n": int(len(hard))}
	torch.save(ck, os.path.join(args.out, "best.pt"))
	print("element head calibration", ck["calibration"], flush=True)
	json.dump({"args": vars(args), "params": nparam, "log": log, "best_score": best}, open(os.path.join(args.out, "train_log.json"), "w"), indent=1)


	if __name__ == "__main__":
	main()