#!/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()