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

Weights, quants, runtime, extension, reports and README swapped to v2.

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+demos/geogebra.jpg filter=lfs diff=lfs merge=lfs -text
+demos/hackernews.jpg filter=lfs diff=lfs merge=lfs -text
+demos/mathworld.jpg filter=lfs diff=lfs merge=lfs -text
+demos/pypi.jpg filter=lfs diff=lfs merge=lfs -text
+demos/python.jpg filter=lfs diff=lfs merge=lfs -text
+demos/wikipedia.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -14,207 +14,321 @@ tags:
 
 # Nadir
 
-**Nadir** is a small on-device model that gives light websites a tasteful, readable dark mode, with the
-care of a real dark theme rather than a filter. It runs entirely in your browser (onnxruntime-web, WASM),
-in about 35 ms of model time per page, and ships with a ready-to-load Chromium extension.
+**Nadir** is a small on-device model that gives light websites a tasteful, readable dark mode, with the care of a
+real dark theme rather than a filter. It runs entirely in your browser (onnxruntime-web, WASM) and ships with a
+ready-to-load Chromium extension. The extension is open source:
+**[github.com/SSHdotCodes/nadir](https://github.com/SSHdotCodes/nadir)**.
+
+Left: the page as published. Right: the same page with Nadir. These are first visits in a fresh profile, rendered by [`training/ext_render.py`](https://huggingface.co/ProCreations/nadir/blob/main/training/ext_render.py).
+
+![geogebra: light page and Nadir dark mode side by side](https://huggingface.co/ProCreations/nadir/resolve/main/demos/geogebra.jpg)
+
+**GeoGebra.** The purple brand color and the illustration keep their colors; surfaces and text go dark.
+
+![python: light page and Nadir dark mode side by side](https://huggingface.co/ProCreations/nadir/resolve/main/demos/python.jpg)
+
+**python.org.** The brand's blue and yellow stay, and the code sample keeps its syntax colors.
+
+![mathworld: light page and Nadir dark mode side by side](https://huggingface.co/ProCreations/nadir/resolve/main/demos/mathworld.jpg)
+
+**Wolfram MathWorld.** Formula images are recolored so they read on the dark page; the rendered sphere keeps its colors.
+
+![hackernews: light page and Nadir dark mode side by side](https://huggingface.co/ProCreations/nadir/resolve/main/demos/hackernews.jpg)
+
+**Hacker News.** The orange header stays; the beige page becomes a warm dark grey.
+
+![addtoany: light page and Nadir dark mode side by side](https://huggingface.co/ProCreations/nadir/resolve/main/demos/addtoany.jpg)
+
+**AddToAny.** The light sunburst texture behind the hero is removed, the dark wordmark turns light, and the share buttons keep their brand colors.
+
+![pypi: light page and Nadir dark mode side by side](https://huggingface.co/ProCreations/nadir/resolve/main/demos/pypi.jpg)
+
+**PyPI.** Badges, code blocks and the blue header keep their colors and hierarchy.
+
+![wikipedia: light page and Nadir dark mode side by side](https://huggingface.co/ProCreations/nadir/resolve/main/demos/wikipedia.jpg)
+
+**Wikipedia** has its own dark theme, so Nadir switches that on instead of designing one.
 
 Nadir follows three rules:
 
 1. **Already dark? Leave it alone.** Pages that are already dark are detected and never touched.
-2. **Has its own dark mode? Use it.** Nadir switches on the site's own dark theme (`prefers-color-scheme`
-   media rules, or `dark`/`theme-dark` class and attribute toggles it detects in the page's CSS).
-3. **No dark mode? The model designs one.** Nadir reads the page's colors and images and predicts a dark
-   palette for that specific site: brand colors stay recognizable, surfaces keep their hierarchy, text stays
-   readable, and logos stay visible: a logo on its own white rectangle is adapted to sit cleanly on the dark
-   surface. Photos, videos and ads are never color-shifted.
-
-In a blind pairwise test on 146 sites it never trained on, judged by Claude Haiku 5.5, the preference
-went to Nadir's dark mode over Dark Reader's on **93 sites vs 50** (3 ties): 65% of the decided pages.
+2. **Has its own dark mode? Use it.** Nadir switches on the site's own dark theme: `prefers-color-scheme` media
+   rules, or `dark`/`theme-dark` class and attribute toggles it detects in the page's CSS.
+3. **No dark mode? The model designs one.** Nadir reads the page's colors and images and designs a dark theme
+   for that specific site:
+   - brand colors stay recognizable and surfaces keep their hierarchy;
+   - text stays readable;
+   - every image, CSS background, sprite and icon is shown the way a designer would choose.
+
+   Photos and videos are never recolored.
+
+**v2 (2026-10-15)** adds an **element head**: the model now decides for itself how to show every image, CSS
+background, sprite icon and `::before` icon, including ones that load after the page. On 23,532 held-out elements
+its choice is acceptable to the teacher 94.4% of the time, against 89.3% for v1's hand-written rules. On 190 held-out live pages, a blind
+Claude Haiku 5.5 judge prefers v2 to Dark Reader 126 to 56 and to v1 91 to 59. Weights, quants and the extension in
+this repository are v2.
+
+## What's new in v2
+
+- **An element head.** v1 decided how to show the first 32 images of a page with a 5-way image head; everything
+  else (CSS backgrounds, sprites, icons, images that load later) fell to hand-written rules, and the rules got logos,
+  icons and light textures wrong often enough to notice. v2 removes those rules. Every visual element is classified by
+  the model itself:
+  - **how to show it:** `keep`, `invert`, `adapt` or `remove`;
+  - **what it is:** one of 9 roles;
+  - **with a calibrated confidence.**
+
+  Three fixed rules remain: photos and people keep their colors, a background is only removed when the head is at
+  least 70% sure, and an element whose surface stays light in the dark design is shown as designed. The element head is a separate 1.2M-parameter graph that reuses the page's design context, so
+  the extension can decide on an image that appears minutes after the page loaded (Gmail icons, lazy-loaded logos,
+  CSS sprites, `::before` icons) in about a millisecond.
+- **Element supervision from Claude Haiku 5.5.** About 80,000 contact sheets were rendered and judged by Haiku 5.5.
+  Each sheet shows elements on their page's real light surface and on the dark surface Haiku designed for that spot,
+  each in every rendering the runtime can produce (kept, inverted, adapted, removed). Haiku picked the best option and
+  marked every acceptable one, which gave 500,785 element judgments. Each option is rendered with the same pixel
+  transforms the extension uses, so what Haiku judged matches what the extension shows.
+- **More pages, and inner pages.** 14,196 new Haiku 5.5 designs, mostly of inner pages: docs, dashboards, inboxes,
+  articles, stores and settings, not just home pages. Together with v1's data that makes 33,615 designs and 6,434
+  native dark-theme pairs.
+- **Runtime fixes found by end-to-end testing.**
+  - Images, backgrounds and icons that appear after the first pass are measured against the page's original light
+    styles.
+  - Long inline (`data:`) images are now measured and classified.
+  - Hotlink-protected images are fetched with the page's origin as the Referer.
+  - A "keep" decision is never overridden by the engine's solid-color heuristic, which misreads SVGs without a
+    fixed size.
+  - SVG logos are recolored shape by shape.
+  - See-through light panels (frosted cards, glass buttons) turn dark.
+  - Opaque surfaces stay opaque. A see-through promo bar would show a hidden header stacked beneath it.
+  - Icons from an external SVG sprite that already follow the page's text color are not inverted a second time.
+  - Logos drawn as CSS masks (the shape painted with a background color) are recolored as ink, like mask icons.
+  - Tabs opened together right after the browser starts no longer miss the model while it is still loading.
 
 ## How it works
 
-Nadir doesn't paint pixels. It redesigns the page's **color system**:
-
-1. **Extraction (`runtime/extract.js`).** One synchronous pass over the DOM collects every distinct color
-   *in its role*: background, text, border, shadow, gradient, SVG fill, SVG stroke. Each token gets 49
-   features: lightness, chroma, hue, area, text size and weight, where it appears (header, nav, link, form,
-   code, footer), what it sits on and the original contrast. It also records four relations between tokens
-   (text-on-surface, surface-on-surface, border-on-surface, same-element) and every visible image or
-   inline-SVG logo with 39 features plus a 32×32 thumbnail. Ad slots are excluded.
-2. **Model (`NadirNet`).** A relation-biased transformer over `[page] + tokens + images`. Each attention head
-   gets a learned bias from the 11 relation channels, so the model knows which text sits on which surface.
-   It predicts for every token an OKLab color as a parameterized edit of the input (lightness, chroma scale,
-   hue rotation, tint, alpha). This keeps brand hues stable and makes impossible colors rare. For every
-   image it predicts a treatment: `keep`, `dim`, `invert`, `light` (monochrome dark glyph → near-white) or
-   `plate` (dark multicolor logo that needs help; shown with a hue-preserving inversion at runtime). A small
-   CNN reads the thumbnails.
-3. **Design rules (`runtime/nadir-infer.js`, `runtime/nadir-map.js`).** Deterministic guards on top of the
-   model:
+Nadir doesn't paint pixels. It redesigns the page's **color system**, and makes a decision for every **visual element**.
+
+1. **Extraction (`runtime/extract.js`).** One synchronous pass over the DOM collects every distinct color *in its role*:
+   background, text, border, shadow, gradient, SVG fill and SVG stroke.
+   - Each color token gets 49 features: lightness, chroma, hue, area, text size and weight, where it appears
+     (header, nav, link, form, code, footer), what it sits on, and the original contrast.
+   - It also records four relations between tokens: text-on-surface, surface-on-surface, border-on-surface and
+     same-element.
+   - Every visual element gets 41 features plus a 32×32 thumbnail. Visual elements are `<img>`, CSS backgrounds,
+     one icon cut out of a sprite sheet, `::before`/`::after` content icons, and SVG images.
+   - Ad slots are excluded.
+2. **Model (`NadirNet`, 16.0M parameters).** A relation-biased transformer over `[page] + colors + images`. Each
+   attention head gets a learned bias from 11 relation channels, so the model knows which text sits on which
+   surface. It has three heads:
+   - **Palette head.** For every color token it predicts an OKLab color as a parameterized edit of the input:
+     lightness, chroma scale, hue rotation, tint and alpha. This keeps brand hues stable and makes impossible
+     colors rare.
+   - **Page context head.** A 32-number summary of the page's design.
+   - **Element head** (new in v2). For any visual element it predicts:
+     - how to show it: `keep`, `invert` (hue-preserving), `adapt` (dark text and lines turn light, brand colors stay,
+       a white box is removed) or `remove` (a light texture behind content);
+     - what it is: photo, illustration, diagram, logo, icon, text/formula, backdrop, avatar or UI decoration.
+
+     It reads the element's thumbnail and features, the page context, the light surface under the element, and the
+     dark surface the palette head designed for that spot. It is exported as its own 1.2M-parameter graph, so the
+     extension can classify elements that appear later (lazy-loaded images, icons in web apps, sprites) without
+     re-running the transformer. Its confidence is calibrated with a fitted temperature.
+3. **Design guards (`runtime/nadir-infer.js`, `runtime/nadir-map.js`).** Deterministic rules on top of the colors:
    - The page canvas is always dark.
-   - Every text and icon color keeps WCAG contrast on *every* surface it actually sits on. It never needs
-     more contrast than the original design had.
+   - Every text and icon color keeps WCAG contrast on *every* surface it actually sits on. It never needs more
+     contrast than the original design had.
+   - A light surface whose text no longer reads is darkened within its own hue.
    - Brand colors keep their hue (±45°).
-   - Things that were invisible in the light design stay invisible: text-fade gradients and borders that
-     matched their surface follow the surface.
-   - Colors drawn over video and photos are left alone when they already read as dark mode (player
-     controls, scrims, captions).
-   - Photos and video are never recolored.
-   - Light, flat background images behind content (paper textures, pale washes, SVG gradients) are cleared so the
-     designed dark surface shows through. Images with real detail are left alone.
-   - Inline SVG logos stay readable: their shapes are drawn as ink, and SVG mask and clip-path shapes (knockout
-     text, cut-outs) are never recolored.
-   - SVG shapes painted in a surface color and invisible on what they sit on (inverted rounded corners, section
-     caps) are part of that surface and take its dark color.
-   - **Logos on their own white rectangle are adapted.** The white background is removed so Nadir's surface
-     shows through, black or gray ink becomes light, and brand colors keep their hue (lifted only as far as
-     needed to read on dark). This applies only to images the page itself marks as its logo.
-   - Dark logos on transparency get a hue-preserving inversion, and already-light ones are kept.
-4. **Engine.** The predicted palette is applied through a patched
-   [Dark Reader](https://github.com/darkreader/darkreader) 4.9.133 dynamic engine (MIT), which handles the
-   hard plumbing: stylesheets, inline styles, shadow DOM and dynamically inserted content. It asks Nadir for
-   every color it rewrites. Colors the model didn't see (new content) are mapped by nearest-neighbor
-   lookup in the predicted palette.
+   - Things that were invisible in the light design stay invisible (text fades, borders and SVG shapes that matched
+     their surface).
+   - Colors over video and photos are left alone when they already read as dark mode.
+
+   For images there are three fixed rules:
+   - Anything the element head itself recognizes as a photo or a person keeps its colors.
+   - A background is only removed when the head is at least 70% sure; otherwise its next choice is used, because
+     taking an image away is the costliest mistake.
+   - An element on a surface that is light in the original and stays light in the dark design (a white carousel
+     button, a kept light card) is shown as designed; recolored, its dark ink would vanish.
+
+   Every other image decision is the model's.
+4. **Engine.** The palette is applied through a patched [Dark Reader](https://github.com/darkreader/darkreader)
+   4.9.133 dynamic engine (MIT). The engine handles stylesheets, inline styles, shadow DOM and dynamically inserted
+   content, and asks Nadir for every color and every background image it rewrites. Colors the model didn't see are
+   mapped by nearest-neighbor lookup in the predicted palette. Background images it didn't see are measured and
+   sent to the element head.
 
 ## Model
 
 | | |
 |---|---|
-| Parameters | 15.9M (15,895,995) |
+| Parameters | 16.0M total, of which the element graph is 1.15M (shared image encoder + element head) |
 | Width / depth / heads | 384 / 8 layers / 6 heads |
-| Inputs | ≤256 color tokens, ≤32 images, 11 relation channels |
-| Outputs | OKLab color + alpha per token, 5-way image treatment |
+| Inputs | ≤256 color tokens, ≤32 images in the page pass, any number of elements in the element pass |
+| Outputs | OKLab color + alpha per token; page context; per element: 4-way treatment + 9-way role (calibrated) |
 | Runtime | onnxruntime-web 1.29 (WASM SIMD, multi-threaded) |
+| Speed | 32 ms per page for the 8-bit page graph (4 WASM threads, Apple M4 Max; 114 ms on one thread), 1.2 s to load it once per browser session; about 0.7 ms per element for the element graph. Revisits use the cached palette. |
 
 ### Files
 
-| File | Precision | Size | Δ color vs fp32 (mean / max, OKLab) | Image agreement |
-|---|---|---|---|---|
-| `onnx/nadir-fp32.onnx` | fp32 | 63.7 MB | 0 / 0 | 100% |
-| `onnx/nadir-int8.onnx` | 8-bit weights | 20.6 MB | 0.00015 / 0.025 | 99.99% |
-| `onnx/nadir-int4.onnx` | 4-bit weights (8-bit input/output layers) | 13.3 MB | 0.00036 / 0.071 | 99.76% |
-
-Fidelity is measured on 400 held-out pages. A difference of 0.02 in OKLab is about the smallest one you can see.
-
-Speed in onnxruntime-web 1.29 (Chrome, Apple M4 Max, 4 WASM threads), per page:
-
-| File | Load | Inference |
+| File | Size | What |
 |---|---|---|
-| fp32 | ~7 s | 34 ms |
-| int8 | ~1.5 s | 33–35 ms |
-| int4 | ~1.0 s | 38 ms |
-
-The extension's own end-to-end time for a first visit is 0.7–3 s, mostly reading the page and its images. Revisits apply
-from cache.
-
-- `onnx/nadir-int8.onnx` is the extension's default.
-- The 8-bit and 4-bit models are **calibrated weight-only quants**:
-  - Layer-sequential GPTQ, with Hessians collected from 256 real crawled pages flowing through the
-    already-quantized earlier layers.
-  - Asymmetric codes in groups of 32, with each group's clipping range chosen by an MSE search.
-  - The 4-bit model keeps its input and output layers at 8 bits.
-- Storage format:
-  - 4-bit layers use ONNX `MatMulNBits` (block 32).
-  - 8-bit layers store uint8 codes with a per-block scale and zero point, dequantized in the graph by standard
-    ops (`Cast`/`Sub`/`Mul`/`Reshape`). onnxruntime folds these into the weights when it loads the model, so
-    the 8-bit layers run as exact fp32 math with no special kernels. Both files are tested on onnxruntime 1.22
-    (CPU) and onnxruntime-web 1.29.
-- Each quant is verified three ways: packing parity against a PyTorch fake-quant model (≤2e-6), fidelity
-  against fp32, and quality against the training targets.
-- They run in onnxruntime-web with outputs matching onnxruntime CPU to ~1e-5.
-- Activation-quantized int8 (ORT dynamic/static) was rejected: worst-case color error up to 0.99.
-
-`model.safetensors` + `model.py` hold the PyTorch weights and architecture.
+| `onnx/nadir-int8.onnx` | 20.6 MB | Page graph, 8-bit. **Shipped in the extension.** Mean color difference to fp32: OKLab ΔE 0.0002 (max 0.024). Its element decisions match fp32 on all 400 validation pages. 32 ms per page in onnxruntime-web (4 threads, M4 Max); loads in 1.2 s. |
+| `onnx/nadir-int4.onnx` | 13.3 MB | Page graph, 4-bit (`MatMulNBits`, blocks of 32; input/output layers stay 8-bit, relation layers fp32). Mean ΔE to fp32: 0.0004 (max 0.11). Element decisions: 99.96% identical. 37 ms per page; loads in 0.9 s. |
+| `onnx/nadir-fp32.onnx` | 63.5 MB | Page graph, full precision. Outputs `lab`, `alpha_logit`, `ctx`. 32 ms per page; loads in 5.6 s. |
+| `onnx/nadir-elem-fp32.onnx` | 4.6 MB | Element graph (1.15M parameters). Inputs `e_f`, `e_thumb`, `under_light`, `under_dark`, `ctx`; outputs `treat_logits` (keep/invert/adapt/remove) and `role_logits` (9 roles). About 0.7 ms per call on a desktop CPU. |
+| `model.safetensors`, `config.json`, `model.py` | 64 MB | PyTorch weights, configuration (including the element-head temperature) and the architecture. |
+| `runtime/` | | The browser runtime: extraction (`extract.js`), decoding and design guards (`nadir-infer.js`), color mapping and image adaptation (`nadir-map.js`), the patched page engine (`darkreader-nadir.js`, MIT). |
+| `training/` | | Crawl, teacher, contact-sheet, labelling, dataset, training, export, quantization and evaluation scripts. |
+| `reports/` | | Training log, export and quantization reports, browser timing (`ortweb_timing_m4max.json`), the held-out element evaluation and the pairwise judge results (`eval_judge.json`). |
+| `nadir-extension.zip` | | The browser extension, ready to load. |
+| `demos/` | | The light vs Nadir comparisons shown above. |
+
+Both quantizations are calibrated on 256 validation pages: per-block scales and zero points are fitted to the
+pages' activations, and the layers that are most sensitive at 4 bits stay at 8 bits or fp32.
 
 ## Training data
 
-- **Pages.** A crawl of 47,569 sites from the [Tranco](https://tranco-list.eu/) top list at 1280×800.
-  Each page was captured in light mode, in emulated dark mode, and with dark-class toggles applied.
-- **Teacher.** 19,563 light pages were redesigned by **Claude Haiku 5.5** (Anthropic Batch API,
-  high effort) acting as the site's own product designer. It was given the screenshot, a contact sheet of
-  the images, and the full color inventory with roles and locations, and returned a dark palette plus image
-  treatments in a strict JSON schema. The design brief is in `training/haiku_gen.py`.
-- **Native pairs.** 4,142 pages with a real built-in dark theme give human-designed light→dark
-  pairs, aligned token by token.
+All teacher data comes from **Claude Haiku 5.5** (Anthropic Batch API, high effort). The total cost for v1 and v2
+together was about $143.
+
+- **Pages.**
+  - v1: a crawl of 47,569 home pages from the [Tranco](https://tranco-list.eu/) top list.
+  - v2 adds 20,375 **inner pages** (docs, dashboards, inboxes, articles, stores, settings) reached from those sites:
+    16,789 light pages, 2,299 with their own dark theme and 271 already dark.
+  - Every page was captured at 1280×800: in light mode, in emulated dark mode, and with dark-class toggles applied.
+- **Palette designs.** Haiku 5.5 redesigned 33,759 light pages (19,563 in v1 and 14,196 new; 33,615 passed QC) as the site's own
+  product designer. It saw the screenshot, a contact sheet of the images, and the full color inventory with roles and
+  locations, and returned a dark palette in a strict JSON schema.
+- **Native pairs.** 6,434 pages with a real built-in dark theme give human-designed light→dark pairs, aligned token by
+  token.
+- **Element judgments (new).** 79,736 contact sheets were rendered by `training/elem_sheets.py`. Each sheet shows
+  elements on their page's real light surface and on the dark surface Haiku designed for that spot. Each element is
+  rendered every way the extension can show it (kept, inverted, adapted, removed), using the extension's own pixel
+  transforms. Haiku 5.5 picked the best rendering, marked every acceptable one and named the element's role:
+  500,785 judgments. On sheets shown to Haiku twice, its two answers agreed 95.5% of the time.
 - **Quality control.**
-  - Designs that missed ids, kept a light canvas, or left >15% of text below 3:1 contrast were dropped.
-  - The rest were passed through the same contrast rule the runtime uses (14,338 of 19,464 designs adjusted).
-  - The train/validation split is by site (hash of the hostname).
-- **Augmentation.** Random token dropout (it also simulates ad blockers and partial loads) and whole-design
-  hue rotation.
+  - Designs that missed ids, kept a light canvas, or left too much text below 3:1 contrast were dropped.
+  - The rest were passed through the same contrast rule as the runtime (25,404 of 40,049 adjusted).
+  - The split is by site (hash of the hostname): 38,076 training pages and 1,973 validation pages. Validation
+    sites hold 23,532 judged elements.
+- **Augmentation.** Random token dropout (it also simulates ad blockers and partial loads) and whole-design hue
+  rotation.
 
 ## Evaluation
 
-The judge was a blind pairwise one: Claude Haiku 5.5 saw the original light page plus two dark versions in
-random order, A and B. It judged readability first, then breakage (glaring light patches, lost logos,
-distorted photos), then design quality. All pages are held-out sites the model never trained on.
+All numbers are on held-out sites that the model never saw in training.
+
+### End to end: pairwise judge
 
-| Nadir vs … (146 held-out sites) | Nadir preferred | Other preferred | Ties |
+Each held-out page was rendered four ways in the same minute, all in a fresh browser profile: as published, with
+Nadir v1, with Nadir v2, and with [Dark Reader](https://github.com/darkreader/darkreader) 4.9 (dynamic mode). Claude
+Haiku 5.5 then saw the light page plus two dark versions in random order, as A and B. It judged readability first,
+then breakage (glaring light patches, lost logos and icons, distorted photos), then design quality.
+
+| Nadir v2 vs … (190 held-out pages) | Nadir v2 preferred | Other preferred | Ties |
 |---|---|---|---|
-| Dark Reader 4.9 (dynamic mode) | **93** | 50 | 3 |
-| Claude Haiku 5.5's own designs (the teacher)¹ | 76 | 61 | 9 |
-| Nadir before its final design rules² | 45 | 29 | 72 |
+| Dark Reader 4.9 (dynamic mode) | **126** | 56 | 8 |
+| Nadir v1 | **91** | 59 | 39 |
+
+v2 is preferred on 61% of the pages with a winner against v1 (91 vs 59) and on 69% against Dark Reader (126 vs 56). Counting only clear preferences (a score gap of at least 2 out of 10), v2 leads v1 39 to 23 and Dark Reader 64 to 27. The pages are live, so re-rendering them changes the counts by about ±5 (carousels, ads and cookie banners differ between renders). Four runtime changes came from reading the judge's reasons on these same pages: opaque surfaces, external SVG sprites, CSS-mask logos and the light-surface rule. Before them, a render batch had v2 and v1 even (77 to 77), so part of that gain may be specific to these pages. Pages that failed to load in any version are left out (189 and 190 of 200).
 
-¹ The teacher's designs were rendered with the earlier runtime, so part of this margin comes from the runtime's newer
-design rules rather than from the model itself. Before those rules, the model and its teacher were even (63 vs 67).
-² The rules for light background textures, SVG logos and SVG masks.
+### Element decisions
 
-**Model size.** A 34.0M-parameter version (512 / 10 layers / 8 heads) trained on the same data was no better: 49 vs 53
-with 45 ties against this one. It reached the same color error and runs about 2× slower in the browser, so the 15.9M
-model is the release.
+For each of 23,532 held-out elements, Haiku 5.5 marked which renderings are acceptable and which one is best. The table
+counts how often the extension's decision is one of the acceptable ones. For v2 that is the element head plus the
+runtime's removal rule: a background is only removed at ≥ 70% confidence. The light-surface rule depends on the rendered
+page, so it is not part of this number. For v1 it is v1's image head plus its
+hand-written rules.
+
+| Elements | n | v1 | **v2** |
+|---|---|---|---|
+| Logos | 3,485 | 80.1% | **93.0%** |
+| Transparent images (icons, line art, formulas) | 7,507 | 87.5% | **93.0%** |
+| Backdrops and textures | 2,111 | 77.7% | **89.6%** |
+| Flat images | 10,181 | 95.9% | **96.8%** |
+| Photos | 148 | 100% | 100% |
+| CSS backgrounds (any type) | 3,952 | 81.7% | **90.3%** |
+| **All** | 23,532 | 89.3% | **94.4%** |
 
-Held-out color error vs the teacher designs (OKLab ΔE): 0.040 (Haiku pages), 0.103
-(native dark themes). Image-treatment accuracy: 79.4% over 5 classes.
+- **Agreement with the teacher.** v2 picks Haiku's first choice 88.6% of the time; v1 did 78.6%.
+- **Calibration.** Confidence is calibrated with a temperature of 1.30 (expected calibration error 0.029).
+  58% of decisions are made with ≥ 90% confidence, and those match Haiku's first choice 98.3% of the time.
+- **Roles.** The element's role (9 classes) is right 78.4% of the time.
+
+### Colors
+
+Held-out color error against the targets (OKLab ΔE, lower is better):
+
+| | Haiku-designed pages | Native dark themes |
+|---|---|---|
+| v2 | **0.0365** | **0.0966** |
+| v1, on its own validation split | 0.0402 | 0.1031 |
+
+**Size.** A 34M-parameter variant was no better in v1's evaluation, so v2 keeps 16.0M parameters. A from-scratch v2
+run (30 epochs) reached 0.0370 / 0.0994 and 94.3% acceptable elements. The released model is warm-started from v1 and
+trained for 18 epochs.
 
 ## Use it
 
 ### Browser extension (Chrome, Edge, Brave, Arc, Aside and other Chromium browsers)
 
-1. Download `nadir-extension.zip` and unzip it.
+1. Download `nadir-extension.zip` (here or from the [GitHub releases](https://github.com/SSHdotCodes/nadir/releases))
+   and unzip it.
 2. Open `chrome://extensions`, enable **Developer mode**, click **Load unpacked**, and choose the folder.
 3. Browse. The toolbar badge shows what Nadir did on each site:
    - `AI`: model-designed dark mode.
    - `ON`: the site's own dark mode.
    - No badge: already dark.
 
-The popup can turn Nadir off for a site.
-
-The extension works alongside content blockers such as uBlock Origin Lite: removed ads simply aren't part
-of the palette, and ads that remain are treated as images (never recolored).
+The popup can turn Nadir off for a site. The extension works alongside content blockers such as uBlock Origin Lite.
 
 ### In your own code
 
 ```js
 import * as ort from 'onnxruntime-web';
 // runtime/extract.js, runtime/nadir-infer.js and runtime/nadir-map.js define the globals used below
-const page = await __nadirExtract({describe: false, thumbs: false});   // in the page: colors, roles, relations, images
+const page = await __nadirExtract({describe: false, thumbs: false, maxImages: 120}); // in the page
 // pixel stats + thumbnails for the images (any context that can fetch them, e.g. an extension page)
 const st = await __nadirImageStats(page.images, {dpr: devicePixelRatio});
 page.images.forEach((im, i) => { Object.assign(im, st[i]); im.f = __nadirImageFeatures(im, page.vw, page.vh, page.docH); });
+const tensor = (data, dims) => new ort.Tensor(data instanceof BigInt64Array ? 'int64' : data instanceof Uint8Array ? 'bool' : 'float32', data, dims);
+const feedsOf = (built) => Object.fromEntries(Object.entries(built.feeds).map(([k, [d, dims]]) => [k, tensor(d, dims)]));
+
+// 1. page graph: the palette and the page context
+const pageSess = await ort.InferenceSession.create('onnx/nadir-int8.onnx', {executionProviders: ['wasm']});
 const built = NadirInfer.buildInputs(page);
-const session = await ort.InferenceSession.create('nadir-int8.onnx', {executionProviders: ['wasm']});
-const feeds = Object.fromEntries(Object.entries(built.feeds).map(([k, [data, dims]]) =>
-  [k, new ort.Tensor(k === 'tok_role' || k === 'src' ? 'int64' : k === 'mask' ? 'bool' : 'float32', data, dims)]));
-const out = await session.run(feeds);
-const {colors, keep, treatments} = NadirInfer.decode(page, built,
-  {lab: out.lab.data, alpha_logit: out.alpha_logit.data, img_logits: out.img_logits.data});
-// colors[i] is the dark replacement for page.tokens[i] (null = keep); treatments[j] applies to page.images[j]
+const out = await pageSess.run(feedsOf(built));
+const {colors, keep} = NadirInfer.decode(page, built, {lab: out.lab.data, alpha_logit: out.alpha_logit.data});
+// colors[i] is the dark replacement for page.tokens[i] (null = keep as is)
+
+// 2. element graph: how to show each image (any time, for any element)
+const elemSess = await ort.InferenceSession.create('onnx/nadir-elem-fp32.onnx', {executionProviders: ['wasm']});
+const {items, idx} = NadirInfer.pageElementItems(page, {lab: out.lab.data});
+const eb = NadirInfer.buildElemInputs(items, out.ctx.data);
+const eo = await elemSess.run(feedsOf(eb));
+const decisions = NadirInfer.decodeElements(eo.treat_logits.data, eo.role_logits.data, eb.K, 1.30); // temperature from config.json
+// decisions[k] = {t: 'keep'|'invert'|'adapt'|'remove', p, role, rp} for page.images[idx[k]]
+// the extension then applies the three fixed image rules above (decide() in the extension's content.js)
 ```
 
 ## Limitations
 
-- Nadir recolors CSS. Text and UI baked into images, canvas drawings and cross-origin iframes keep their
-  original colors, though the iframes get their own Nadir pass in the extension.
-- Colors are mapped per (color, role). If one exact color is used both as a light card and as a light
-  button, both get the same dark color.
-- Background images that mix light areas with strong color (for example a colored hero shape with a white base) are
-  kept as they are, so their light parts stay light.
-- The model reads at most 256 colors per page. On very colorful pages the rarest colors are mapped by nearest
-  neighbor in the predicted palette.
-- The teacher is an LLM. Its taste is good but not perfect, and the model inherits its preferences.
-- The judge is also an LLM, so the evaluation is automated, not a human study.
+- Nadir recolors CSS and decides per image; it never repaints pixels inside an image except when adapting one (dark
+  ink → light). Text baked into photos, canvas drawings and cross-origin iframes keep their colors; the iframes get
+  their own Nadir pass in the extension.
+- Colors are mapped per (color, role). If one exact color is used both as a light card and as a light button, both
+  get the same dark color.
+- The page pass reads at most 256 colors. On very colorful pages the rarest colors are mapped by nearest neighbor.
+- Glyphs drawn with CSS instead of images (hamburger bars made of thin boxes, chevrons made of borders, checkmarks)
+  follow the per-color mapping, so a dark bar can stay dark on the dark page. White glyphs inside colored SVG icon
+  buttons can turn dark; the brand color stays, so they still read, except on black buttons.
+- A sprite sheet drawn as a CSS background gets one decision for all its icons (the majority of the icons the
+  element head saw), because the engine rewrites the CSS rule, not each element.
+- Teacher and judge are both LLMs (Claude Haiku 5.5). Their taste is good but not perfect, the model inherits it,
+  and the evaluation is automated, not a human study.
 
 ## License
 
 Apache-2.0 for the model, weights and Nadir code. The bundled page engine is a modified
-[Dark Reader](https://github.com/darkreader/darkreader) (MIT License, © Dark Reader Ltd.); its license is
-included in the extension (`lib/DARKREADER-LICENSE`).
+[Dark Reader](https://github.com/darkreader/darkreader) (MIT License, © Dark Reader Ltd.); its license is included in
+the extension (`lib/DARKREADER-LICENSE`).

config.json

@@ -5,32 +5,66 @@
  "ff": 4,
  "drop": 0.1,
  "architecture": "NadirNet",
+ "version": 2,
  "token_features": 49,
- "image_features": 39,
+ "image_features": 41,
  "relation_channels": 11,
  "max_tokens": 256,
  "max_images": 32,
- "image_treatments": [
+ "page_graph": {
+  "inputs": [
+   "tok_f",
+   "tok_role",
+   "tok_lab",
+   "img_f",
+   "img_thumb",
+   "page_f",
+   "src",
+   "rel",
+   "mask"
+  ],
+  "outputs": [
+   "lab",
+   "alpha_logit",
+   "ctx"
+  ]
+ },
+ "element_graph": {
+  "inputs": [
+   "e_f",
+   "e_thumb",
+   "under_light",
+   "under_dark",
+   "ctx"
+  ],
+  "outputs": [
+   "treat_logits",
+   "role_logits"
+  ]
+ },
+ "element_treatments": [
   "keep",
-  "dim",
   "invert",
-  "light",
-  "plate"
+  "adapt",
+  "remove"
  ],
- "inputs": [
-  "tok_f",
-  "tok_role",
-  "tok_lab",
-  "img_f",
-  "img_thumb",
-  "page_f",
-  "src",
-  "rel",
-  "mask"
+ "image_roles": [
+  "photo",
+  "illustration",
+  "diagram",
+  "logo",
+  "icon",
+  "text",
+  "backdrop",
+  "avatar",
+  "ui"
  ],
- "outputs": [
-  "lab",
-  "alpha_logit",
-  "img_logits"
- ]
+ "context_dim": 32,
+ "element_temperature": 1.3000000000000007,
+ "element_calibration": {
+  "temperature": 1.3000000000000007,
+  "nll": 0.32924941182136536,
+  "ece": 0.028595900163054466,
+  "n": 23532
+ }
 }

model.py

@@ -5,15 +5,24 @@ Attention logits get a learned per-head bias from the page's structure graph (te
 surface-on-surface, border-on-surface, same-element, image-on-surface, same-colour), so the model
 reasons about hierarchy and contrast the way a designer would. Outputs a dark colour per token
 (OKLab, parameterised relative to the input so hue is kept unless the model chooses otherwise),
-an alpha, and a treatment class per image.
+an alpha, and a compact page context.
+
+Element head (v2): decides how every image, CSS background, sprite or pseudo-element icon on the page is shown
+(keep / invert / adapt / remove) and what it is (photo, logo, icon, formula...). It reuses the trunk's image
+encoder and reads the page context plus the light surface under the element and the dark surface the palette
+head designed for it, so it can run for any element at any time (late-loaded icons, sprites) in a tiny second
+graph without re-running the transformer.
 """
 import math
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-FT, FI, FP = 49, 39, 8
-N_ROLES, N_TREAT = 7, 5
+FT, FI, FP = 49, 41, 8
+N_ROLES = 7
+TREAT = ["keep", "invert", "adapt", "remove"]
+IMG_ROLES = ["photo", "illustration", "diagram", "logo", "icon", "text", "backdrop", "avatar", "ui"]
+N_TREAT, N_IROLE, N_CTX = len(TREAT), len(IMG_ROLES), 32
 K_REL = 11
 
 
@@ -89,7 +98,19 @@ class NadirNet(nn.Module):
         self.blocks = nn.ModuleList([Block(d, heads, ff, drop) for _ in range(layers)])
         self.ln = nn.LayerNorm(d)
         self.color_head = MLP(d + 3, d, 6)
-        self.img_head = MLP(d, d // 2, N_TREAT)
+        self.ctx_head = nn.Linear(d, N_CTX)
+        self.elem_head = MLP(d + N_CTX + 72, d, N_TREAT + N_IROLE, drop=drop)
+
+    def embed_images(self, img_f, img_thumb):
+        return self.img_in(torch.cat([img_f, self.cnn(img_thumb)], -1))
+
+    def elements(self, e_f, e_thumb, under_light, under_dark, ctx):
+        """e_f (B,K,FI), e_thumb (B,K,4,32,32), under_light/under_dark (B,K,3) OKLab, ctx (B,N_CTX)
+        -> treatment logits (B,K,N_TREAT), role logits (B,K,N_IROLE)."""
+        K = e_f.shape[1]
+        h = torch.cat([self.embed_images(e_f, e_thumb), ctx[:, None].expand(-1, K, -1), fourier_lab(under_light), fourier_lab(under_dark)], -1)
+        z = self.elem_head(h)
+        return z[..., :N_TREAT], z[..., N_TREAT:]
 
     def forward(self, tok_f, tok_role, tok_lab, img_f, img_thumb, page_f, src, rel, mask):
         """tok_f (B,N,FT) tok_role (B,N) long, tok_lab (B,N,3), img_f (B,M,FI), img_thumb (B,M,4,32,32) float 0-1,
@@ -98,7 +119,7 @@ class NadirNet(nn.Module):
         M = img_f.shape[1]
         tok = self.tok_in(torch.cat([tok_f, F.one_hot(tok_role, N_ROLES).to(tok_f.dtype), fourier_lab(tok_lab)], -1)) + self.typ.weight[1]
         if M > 0:
-            img = self.img_in(torch.cat([img_f, self.cnn(img_thumb)], -1)) + self.typ.weight[2]
+            img = self.embed_images(img_f, img_thumb) + self.typ.weight[2]
         else:
             img = tok.new_zeros(B, 0, tok.shape[-1])
         cls = self.page_in(page_f)[:, None] + self.typ.weight[0]
@@ -117,8 +138,18 @@ class NadirNet(nn.Module):
         b = r * a_in + s * b_in + 0.12 * torch.tanh(z[..., 4])
         lab = torch.stack([L, a, b], -1)
         alpha_logit = z[..., 5]
-        img_logits = self.img_head(x[:, 1 + N:]) if M > 0 else x.new_zeros(B, 0, N_TREAT)
-        return lab, alpha_logit, img_logits
+        ctx = self.ctx_head(x[:, 0])
+        return lab, alpha_logit, ctx
+
+
+class ElemNet(nn.Module):
+    """The element head as its own graph (shares NadirNet's weights): classify any set of page images."""
+    def __init__(self, net):
+        super().__init__()
+        self.net = net
+
+    def forward(self, e_f, e_thumb, under_light, under_dark, ctx):
+        return self.net.elements(e_f, e_thumb, under_light, under_dark, ctx)
 
 
 # ---------- colour maths (differentiable) ----------

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e7acc89ba8a039a4889d89ec9f01afce49b2f765d1ac57a489eb4c604aef2ceb
-size 63596788
+oid sha256:c6a783d4cafe79ea3ddae0cfa1bda3b9f378d8298ba98551863ad84f48134fad
+size 64123972

nadir-extension.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:71e3cb4fe7f3925621abf48d2230d97d18622da6cf08dcea42a8b5959acc0219
-size 23422165
+oid sha256:3a7da36bb63e9fa6967c41dde36d87f5fb4ff535fb77dfbda87ad2831e53adc7
+size 27628972

onnx/nadir-elem-fp32.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:57f4e5e59a5cf9a45dc7c95cc38a88194469554d157dba7dc821567f5d4863fa
+size 4615034

onnx/nadir-fp32.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:83c54558313415596d4d3b3e753ff30e3a72cdc65e27a5c8eefcb132b9d5c934
-size 63699911
+oid sha256:b729c8633572ce062616cecc3ae0d2fd416e797b39f384a631963c8853755c75
+size 63453840

onnx/nadir-int4.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5e11864ad3d7f1edb7bc8b7006e5d7a1c12401bbc17578f590182c0ca764802a
-size 13313059
+oid sha256:d404b6eb6eb5cc494e8bef75486d34947a6b981776b4c2743bba765ec17e99a5
+size 13271851

onnx/nadir-int8.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:bbdb0a2b13bfca0c05f8d19d54d3b721fb4fcd459d40e8519a7a5f77af15d634
-size 20626640
+oid sha256:3db9718c191de95a527ff160f798e24401981236120019842ecc4d72022acee8
+size 20585432

reports/elem_eval.json

@@ -0,0 +1,171 @@
+{
+ "temperature": 1.3000000000000007,
+ "overall": {
+  "v2": {
+   "n": 23532,
+   "acc": 0.8836,
+   "acceptable": 0.9422
+  },
+  "v1": {
+   "n": 23532,
+   "acc": 0.7859,
+   "acceptable": 0.8925
+  }
+ },
+ "by_bucket": {
+  "backdrop": {
+   "v2": {
+    "n": 2111,
+    "acc": 0.8015,
+    "acceptable": 0.8825
+   },
+   "v1": {
+    "n": 2111,
+    "acc": 0.6902,
+    "acceptable": 0.7774
+   }
+  },
+  "flat": {
+   "v2": {
+    "n": 10181,
+    "acc": 0.9376,
+    "acceptable": 0.9671
+   },
+   "v1": {
+    "n": 10181,
+    "acc": 0.9071,
+    "acceptable": 0.9587
+   }
+  },
+  "icon": {
+   "v2": {
+    "n": 100,
+    "acc": 0.79,
+    "acceptable": 0.91
+   },
+   "v1": {
+    "n": 100,
+    "acc": 0.69,
+    "acceptable": 0.91
+   }
+  },
+  "logo": {
+   "v2": {
+    "n": 3485,
+    "acc": 0.8482,
+    "acceptable": 0.9286
+   },
+   "v1": {
+    "n": 3485,
+    "acc": 0.7228,
+    "acceptable": 0.8009
+   }
+  },
+  "photo": {
+   "v2": {
+    "n": 148,
+    "acc": 1.0,
+    "acceptable": 1.0
+   },
+   "v1": {
+    "n": 148,
+    "acc": 1.0,
+    "acceptable": 1.0
+   }
+  },
+  "transparent": {
+   "v2": {
+    "n": 7507,
+    "acc": 0.8489,
+    "acceptable": 0.9307
+   },
+   "v1": {
+    "n": 7507,
+    "acc": 0.6748,
+    "acceptable": 0.8753
+   }
+  }
+ },
+ "by_kind": {
+  "bg": {
+   "v2": {
+    "n": 3952,
+    "acc": 0.8072,
+    "acceptable": 0.8957
+   },
+   "v1": {
+    "n": 3952,
+    "acc": 0.684,
+    "acceptable": 0.8171
+   }
+  },
+  "img": {
+   "v2": {
+    "n": 19534,
+    "acc": 0.8992,
+    "acceptable": 0.9517
+   },
+   "v1": {
+    "n": 19534,
+    "acc": 0.8069,
+    "acceptable": 0.9078
+   }
+  },
+  "pseudo": {
+   "v2": {
+    "n": 13,
+    "acc": 0.9231,
+    "acceptable": 0.9231
+   },
+   "v1": {
+    "n": 13,
+    "acc": 0.7692,
+    "acceptable": 0.9231
+   }
+  },
+  "svgimage": {
+   "v2": {
+    "n": 33,
+    "acc": 0.8182,
+    "acceptable": 0.8788
+   },
+   "v1": {
+    "n": 33,
+    "acc": 0.5758,
+    "acceptable": 0.8788
+   }
+  }
+ },
+ "calibration": {
+  "0.3": {
+   "n": 88,
+   "acc": 0.477
+  },
+  "0.4": {
+   "n": 599,
+   "acc": 0.422
+  },
+  "0.5": {
+   "n": 1989,
+   "acc": 0.598
+  },
+  "0.6": {
+   "n": 2682,
+   "acc": 0.76
+  },
+  "0.7": {
+   "n": 2175,
+   "acc": 0.812
+  },
+  "0.8": {
+   "n": 2292,
+   "acc": 0.884
+  },
+  "0.9": {
+   "n": 13707,
+   "acc": 0.983
+  }
+ },
+ "role_acc": 0.7843,
+ "n_double_labelled": 657
+}

reports/export_report.json

@@ -1,20 +1,42 @@
 {
- "params": 15895995,
- "parity_max_abs": 1.2516975402832031e-06,
+ "params": 16027747,
+ "elem_params": 1147181,
+ "parity_max_abs": 2.7418136596679688e-06,
+ "elem_parity_max_abs": 1.9073486328125e-05,
+ "temperature": 1.3000000000000007,
+ "treat": [
+  "keep",
+  "invert",
+  "adapt",
+  "remove"
+ ],
+ "img_roles": [
+  "photo",
+  "illustration",
+  "diagram",
+  "logo",
+  "icon",
+  "text",
+  "backdrop",
+  "avatar",
+  "ui"
+ ],
+ "elem_val_sample": {
+  "n": 859,
+  "acc": 0.9068684516880093,
+  "ok": 0.9522700814901047
+ },
  "variants": {
   "fp32": {
    "dE_vs_fp32": 0.0,
    "dE_max": 0.0,
-   "img_agree": 1.0,
-   "size_mb": 63.7,
-   "cpu_ms_median": 7.48
-  },
-  "fp16": {
-   "dE_vs_fp32": 0.0002421401860071152,
-   "dE_max": 0.23341728746891022,
-   "img_agree": 1.0,
-   "size_mb": 31.91,
-   "cpu_ms_median": 10.17
+   "ctx_max_abs": 0.0,
+   "size_mb": 63.45,
+   "cpu_ms_median": 10.48
   }
+ },
+ "elem": {
+  "size_mb": 4.62,
+  "cpu_ms_median": 0.65
  }
 }

reports/ortweb_timing_m4max.json

@@ -1,67 +1,57 @@
 {
- "runtime": "onnxruntime-web 1.29 wasm, Chrome for Testing 141 headless, Apple M4 Max; t1/t4 = 1/4 threads; 12 held-out pages, median of 3 runs each",
+ "runtime": "onnxruntime-web 1.29 wasm, Chrome for Testing 141 headless, Apple M4 Max; t1/t4 = 1/4 threads; 12 held-out pages, median of 3 runs each; outputs compared with onnxruntime CPU",
  "fp32": {
   "t1": {
    "isolated": false,
-   "load_ms": 7246.9,
-   "ms_median": 117.08,
-   "max_abs_lab_vs_ort_cpu": 6.16908073425293e-06,
-   "max_abs_img_vs_ort_cpu": 4.172325134277344e-06
+   "load_ms": 5602.2,
+   "ms_median": 112.13,
+   "max_abs_lab_vs_ort_cpu": 3.3974647521972656e-06,
+   "max_abs_alpha_logit_vs_ort_cpu": 1.4543533325195312e-05,
+   "max_abs_ctx_vs_ort_cpu": 3.5762786865234375e-06
   },
   "t4": {
    "isolated": false,
-   "load_ms": 7363.2,
-   "ms_median": 33.93,
-   "max_abs_lab_vs_ort_cpu": 6.16908073425293e-06,
-   "max_abs_img_vs_ort_cpu": 4.172325134277344e-06
+   "load_ms": 5554.2,
+   "ms_median": 31.8,
+   "max_abs_lab_vs_ort_cpu": 3.3974647521972656e-06,
+   "max_abs_alpha_logit_vs_ort_cpu": 1.4543533325195312e-05,
+   "max_abs_ctx_vs_ort_cpu": 3.5762786865234375e-06
   }
  },
  "int8": {
   "t1": {
    "isolated": false,
-   "load_ms": 1356.6,
-   "ms_median": 116.78,
-   "max_abs_lab_vs_ort_cpu": 1.9669532775878906e-06,
-   "max_abs_img_vs_ort_cpu": 2.9802322387695312e-06
+   "load_ms": 1246,
+   "ms_median": 113.62,
+   "max_abs_lab_vs_ort_cpu": 4.0531158447265625e-06,
+   "max_abs_alpha_logit_vs_ort_cpu": 2.86102294921875e-05,
+   "max_abs_ctx_vs_ort_cpu": 4.0531158447265625e-06
   },
   "t4": {
    "isolated": false,
-   "load_ms": 1628.3,
-   "ms_median": 34.78,
-   "max_abs_lab_vs_ort_cpu": 1.9669532775878906e-06,
-   "max_abs_img_vs_ort_cpu": 2.9802322387695312e-06
-  }
- },
- "int8_repeat": {
-  "t1": {
-   "isolated": false,
-   "load_ms": 1523.5,
-   "ms_median": 115.82,
-   "max_abs_lab_vs_ort_cpu": 1.9669532775878906e-06,
-   "max_abs_img_vs_ort_cpu": 2.9802322387695312e-06
-  },
-  "t4": {
-   "isolated": false,
-   "load_ms": 1519,
-   "ms_median": 32.68,
-   "max_abs_lab_vs_ort_cpu": 1.9669532775878906e-06,
-   "max_abs_img_vs_ort_cpu": 2.9802322387695312e-06
+   "load_ms": 1214.3,
+   "ms_median": 32.03,
+   "max_abs_lab_vs_ort_cpu": 4.0531158447265625e-06,
+   "max_abs_alpha_logit_vs_ort_cpu": 2.86102294921875e-05,
+   "max_abs_ctx_vs_ort_cpu": 4.0531158447265625e-06
   }
  },
  "int4": {
   "t1": {
    "isolated": false,
-   "load_ms": 1074.8,
-   "ms_median": 137.33,
-   "max_abs_lab_vs_ort_cpu": 1.3709068298339844e-06,
-   "max_abs_img_vs_ort_cpu": 7.659196853637695e-06
+   "load_ms": 854.5,
+   "ms_median": 132.23,
+   "max_abs_lab_vs_ort_cpu": 6.377696990966797e-06,
+   "max_abs_alpha_logit_vs_ort_cpu": 3.719329833984375e-05,
+   "max_abs_ctx_vs_ort_cpu": 3.337860107421875e-06
   },
   "t4": {
    "isolated": false,
-   "load_ms": 1047,
-   "ms_median": 38.47,
-   "max_abs_lab_vs_ort_cpu": 1.3709068298339844e-06,
-   "max_abs_img_vs_ort_cpu": 7.659196853637695e-06
+   "load_ms": 857.2,
+   "ms_median": 37.18,
+   "max_abs_lab_vs_ort_cpu": 6.377696990966797e-06,
+   "max_abs_alpha_logit_vs_ort_cpu": 3.719329833984375e-05,
+   "max_abs_ctx_vs_ort_cpu": 3.337860107421875e-06
   }
  }
 }

reports/quant_report.json

@@ -17,7 +17,8 @@
   "blocks.4.rel",
   "blocks.5.rel",
   "blocks.6.rel",
-  "blocks.7.rel"
+  "blocks.7.rel",
+  "ctx_head"
  ],
  "int4_8bit_edge_layers": [
   "tok_in.net.0",
@@ -25,39 +26,36 @@
   "img_in.net.0",
   "img_in.net.3",
   "color_head.net.0",
-  "color_head.net.3",
-  "img_head.net.0",
-  "img_head.net.3"
+  "color_head.net.3"
  ],
  "variants": {
   "fp32": {
-   "haiku_dE": 0.039349029815445344,
-   "native_dE": 0.08639045081450604,
-   "img_acc": 0.7966292134831461,
-   "size_mb": 63.7,
-   "cpu_ms_median": 11.42
+   "haiku_dE": 0.03649106619724383,
+   "native_dE": 0.07405353200621903,
+   "size_mb": 63.45,
+   "cpu_ms_median": 9.31
   },
   "int8": {
-   "dE_vs_fp32": 0.00015194275738394936,
-   "dE_max": 0.02521813102066517,
-   "img_agree": 0.9998743718592965,
-   "haiku_dE": 0.03935358844697476,
-   "native_dE": 0.0863548206735868,
-   "img_acc": 0.7968164794007491,
-   "pack_parity_max_abs": 2.5033950805664062e-06,
-   "size_mb": 20.63,
-   "cpu_ms_median": 9.19
+   "dE_vs_fp32": 0.0001628809163776168,
+   "dE_max": 0.02378411777317524,
+   "ctx_max_abs": 0.013257980346679688,
+   "haiku_dE": 0.036490150559693575,
+   "native_dE": 0.074042251650244,
+   "elem_decisions_same_as_fp32": 1.0,
+   "pack_parity_max_abs": 1.7881393432617188e-06,
+   "size_mb": 20.59,
+   "cpu_ms_median": 9.77
   },
   "int4": {
-   "dE_vs_fp32": 0.00036296586506068707,
-   "dE_max": 0.0709674283862114,
-   "img_agree": 0.9976130653266332,
-   "haiku_dE": 0.039340872559696434,
-   "native_dE": 0.08632776886341162,
-   "img_acc": 0.797003745318352,
-   "pack_parity_max_abs": 1.5497207641601562e-06,
-   "size_mb": 13.31,
-   "cpu_ms_median": 9.76
+   "dE_vs_fp32": 0.00039874050049547804,
+   "dE_max": 0.10922706872224808,
+   "ctx_max_abs": 0.07820296287536621,
+   "haiku_dE": 0.03647379979180793,
+   "native_dE": 0.07408187076449395,
+   "elem_decisions_same_as_fp32": 0.9995768986672308,
+   "pack_parity_max_abs": 2.9206275939941406e-06,
+   "size_mb": 13.27,
+   "cpu_ms_median": 8.29
   }
  }
 }

reports/train_log.json

@@ -1,501 +1,367 @@
 {
  "args": {
-  "data": "/home/user/.local/share/rtx-temporary-jobs/tj-nadir-0fe76575c0f3/payload/work/data/ds_final.pkl",
-  "out": "/job/export/run_m",
+  "data": "ds_v2.pkl",
+  "out": "run_w",
   "d": 384,
   "layers": 8,
   "heads": 6,
-  "epochs": 30,
+  "epochs": 18,
   "bs": 48,
-  "lr": 0.0004,
+  "lr": 0.0003,
   "wd": 0.05,
   "drop": 0.1,
   "native_weight": 0.5,
   "seed": 0,
-  "workers": 4
+  "workers": 5,
+  "elem_weight": 0.1,
+  "role_weight": 0.03,
+  "elem_balance": 0.3,
+  "init": "v1_best.pt"
  },
- "params": 15895995,
+ "params": 16027747,
  "log": [
   {
    "epoch": 0,
-   "step": 468,
-   "secs": 81,
-   "loss": 0.3405,
-   "dE": 0.12132,
-   "con": 0.50822,
-   "img": 1.29224,
-   "s0_dE": 0.06306,
-   "s0_dL": 0.06024,
-   "s0_con": 0.04883,
-   "img_acc": 0.73416,
-   "s1_dE": 0.16366,
-   "s1_dL": 0.15909,
-   "s1_con": 0.11938
+   "step": 794,
+   "secs": 126,
+   "loss": 0.1623,
+   "dE": 0.04288,
+   "con": 0.03147,
+   "elem": 0.37817,
+   "e_ok": 0.94195,
+   "s0_dE": 0.04166,
+   "s0_dL": 0.03983,
+   "s0_con": 0.01493,
+   "e_acc": 0.87222,
+   "e_nll": 0.31791,
+   "r_acc": 0.74936,
+   "s1_dE": 0.11709,
+   "s1_dL": 0.11392,
+   "s1_con": 0.29847
   },
   {
    "epoch": 1,
-   "step": 936,
-   "secs": 148,
-   "loss": 0.18873,
-   "dE": 0.06161,
-   "con": 0.03771,
-   "img": 0.99138,
-   "s0_dE": 0.05432,
-   "s0_dL": 0.05155,
-   "s0_con": 0.03668,
-   "img_acc": 0.78603,
-   "s1_dE": 0.1523,
-   "s1_dL": 0.14838,
-   "s1_con": 0.08049
+   "step": 1588,
+   "secs": 246,
+   "loss": 0.14616,
+   "dE": 0.044,
+   "con": 0.03079,
+   "elem": 0.28204,
+   "e_ok": 0.94722,
+   "s0_dE": 0.04111,
+   "s0_dL": 0.03933,
+   "s0_con": 0.01486,
+   "e_acc": 0.88492,
+   "e_nll": 0.29707,
+   "r_acc": 0.76823,
+   "s1_dE": 0.11303,
+   "s1_dL": 0.10971,
+   "s1_con": 0.29463
   },
   {
    "epoch": 2,
-   "step": 1404,
-   "secs": 215,
-   "loss": 0.17375,
-   "dE": 0.05668,
-   "con": 0.03498,
-   "img": 0.90371,
-   "s0_dE": 0.04977,
-   "s0_dL": 0.04725,
-   "s0_con": 0.02243,
-   "img_acc": 0.78005,
-   "s1_dE": 0.14476,
-   "s1_dL": 0.14153,
-   "s1_con": 0.26806
+   "step": 2382,
+   "secs": 371,
+   "loss": 0.14248,
+   "dE": 0.0433,
+   "con": 0.02879,
+   "elem": 0.2718,
+   "e_ok": 0.94794,
+   "s0_dE": 0.04045,
+   "s0_dL": 0.03863,
+   "s0_con": 0.01225,
+   "e_acc": 0.88913,
+   "e_nll": 0.28905,
+   "r_acc": 0.77894,
+   "s1_dE": 0.11051,
+   "s1_dL": 0.1072,
+   "s1_con": 0.26784
   },
   {
    "epoch": 3,
-   "step": 1872,
-   "secs": 281,
-   "loss": 0.16779,
-   "dE": 0.05488,
-   "con": 0.04139,
-   "img": 0.86775,
-   "s0_dE": 0.0506,
-   "s0_dL": 0.04808,
-   "s0_con": 0.0331,
-   "img_acc": 0.79611,
-   "s1_dE": 0.13697,
-   "s1_dL": 0.13372,
-   "s1_con": 0.1237
+   "step": 3176,
+   "secs": 481,
+   "loss": 0.13929,
+   "dE": 0.04303,
+   "con": 0.0288,
+   "elem": 0.26188,
+   "e_ok": 0.94794,
+   "s0_dE": 0.04085,
+   "s0_dL": 0.03902,
+   "s0_con": 0.01576,
+   "e_acc": 0.88403,
+   "e_nll": 0.29122,
+   "r_acc": 0.78404,
+   "s1_dE": 0.12151,
+   "s1_dL": 0.118,
+   "s1_con": 0.13395
   },
   {
    "epoch": 4,
-   "step": 2340,
-   "secs": 349,
-   "loss": 0.16278,
-   "dE": 0.05274,
-   "con": 0.0414,
-   "img": 0.85412,
-   "s0_dE": 0.04849,
-   "s0_dL": 0.04604,
-   "s0_con": 0.02439,
-   "img_acc": 0.78982,
-   "s1_dE": 0.13195,
-   "s1_dL": 0.12841,
-   "s1_con": 0.16648
+   "step": 3970,
+   "secs": 592,
+   "loss": 0.13736,
+   "dE": 0.04279,
+   "con": 0.029,
+   "elem": 0.25577,
+   "e_ok": 0.94722,
+   "s0_dE": 0.0397,
+   "s0_dL": 0.03794,
+   "s0_con": 0.01431,
+   "e_acc": 0.88994,
+   "e_nll": 0.29091,
+   "r_acc": 0.78723,
+   "s1_dE": 0.11205,
+   "s1_dL": 0.10875,
+   "s1_con": 0.24692
   },
   {
    "epoch": 5,
-   "step": 2808,
-   "secs": 416,
-   "loss": 0.15903,
-   "dE": 0.05166,
-   "con": 0.03833,
-   "img": 0.84086,
-   "s0_dE": 0.0472,
-   "s0_dL": 0.04477,
-   "s0_con": 0.03465,
-   "img_acc": 0.80467,
-   "s1_dE": 0.12882,
-   "s1_dL": 0.12538,
-   "s1_con": 0.12963
+   "step": 4764,
+   "secs": 705,
+   "loss": 0.13486,
+   "dE": 0.04215,
+   "con": 0.0303,
+   "elem": 0.24835,
+   "e_ok": 0.95037,
+   "s0_dE": 0.03995,
+   "s0_dL": 0.03818,
+   "s0_con": 0.01163,
+   "e_acc": 0.89342,
+   "e_nll": 0.2904,
+   "r_acc": 0.78489,
+   "s1_dE": 0.10674,
+   "s1_dL": 0.1034,
+   "s1_con": 0.21942
   },
   {
    "epoch": 6,
-   "step": 3276,
-   "secs": 484,
-   "loss": 0.15729,
-   "dE": 0.05076,
-   "con": 0.04035,
-   "img": 0.83185,
-   "s0_dE": 0.0457,
-   "s0_dL": 0.04338,
-   "s0_con": 0.02999,
-   "img_acc": 0.80345,
-   "s1_dE": 0.15936,
-   "s1_dL": 0.15619,
-   "s1_con": 0.51491
+   "step": 5558,
+   "secs": 816,
+   "loss": 0.13262,
+   "dE": 0.04164,
+   "con": 0.02965,
+   "elem": 0.24241,
+   "e_ok": 0.94998,
+   "s0_dE": 0.03888,
+   "s0_dL": 0.03715,
+   "s0_con": 0.01823,
+   "e_acc": 0.89142,
+   "e_nll": 0.29542,
+   "r_acc": 0.78795,
+   "s1_dE": 0.10818,
+   "s1_dL": 0.10477,
+   "s1_con": 0.13355
   },
   {
    "epoch": 7,
-   "step": 3744,
-   "secs": 552,
-   "loss": 0.15392,
-   "dE": 0.04978,
-   "con": 0.0403,
-   "img": 0.81841,
-   "s0_dE": 0.04612,
-   "s0_dL": 0.04395,
-   "s0_con": 0.01962,
-   "img_acc": 0.79899,
-   "s1_dE": 0.14854,
-   "s1_dL": 0.14578,
-   "s1_con": 0.42056
+   "step": 6352,
+   "secs": 927,
+   "loss": 0.1297,
+   "dE": 0.04077,
+   "con": 0.02959,
+   "elem": 0.23773,
+   "e_ok": 0.94318,
+   "s0_dE": 0.03864,
+   "s0_dL": 0.03688,
+   "s0_con": 0.01412,
+   "e_acc": 0.88339,
+   "e_nll": 0.29541,
+   "r_acc": 0.7865,
+   "s1_dE": 0.10767,
+   "s1_dL": 0.10424,
+   "s1_con": 0.15202
   },
   {
    "epoch": 8,
-   "step": 4212,
-   "secs": 619,
-   "loss": 0.15244,
-   "dE": 0.04894,
-   "con": 0.04062,
-   "img": 0.81195,
-   "s0_dE": 0.04632,
-   "s0_dL": 0.04425,
-   "s0_con": 0.02413,
-   "img_acc": 0.79917,
-   "s1_dE": 0.12446,
-   "s1_dL": 0.12131,
-   "s1_con": 0.23072
+   "step": 7146,
+   "secs": 1038,
+   "loss": 0.12712,
+   "dE": 0.0403,
+   "con": 0.02936,
+   "elem": 0.22914,
+   "e_ok": 0.94752,
+   "s0_dE": 0.03815,
+   "s0_dL": 0.03647,
+   "s0_con": 0.01527,
+   "e_acc": 0.89185,
+   "e_nll": 0.30129,
+   "r_acc": 0.78701,
+   "s1_dE": 0.10567,
+   "s1_dL": 0.10227,
+   "s1_con": 0.20398
   },
   {
    "epoch": 9,
-   "step": 4680,
-   "secs": 687,
-   "loss": 0.14931,
-   "dE": 0.0481,
-   "con": 0.03785,
-   "img": 0.79837,
-   "s0_dE": 0.04452,
-   "s0_dL": 0.04238,
-   "s0_con": 0.02434,
-   "img_acc": 0.80094,
-   "s1_dE": 0.12096,
-   "s1_dL": 0.11779,
-   "s1_con": 0.21844
+   "step": 7940,
+   "secs": 1148,
+   "loss": 0.12518,
+   "dE": 0.0398,
+   "con": 0.0268,
+   "elem": 0.22563,
+   "e_ok": 0.94726,
+   "s0_dE": 0.03823,
+   "s0_dL": 0.03653,
+   "s0_con": 0.01809,
+   "e_acc": 0.88853,
+   "e_nll": 0.30672,
+   "r_acc": 0.78863,
+   "s1_dE": 0.10322,
+   "s1_dL": 0.09977,
+   "s1_con": 0.15727
   },
   {
    "epoch": 10,
-   "step": 5148,
-   "secs": 755,
-   "loss": 0.1484,
-   "dE": 0.04786,
-   "con": 0.03854,
-   "img": 0.79393,
-   "s0_dE": 0.04549,
-   "s0_dL": 0.04344,
-   "s0_con": 0.02666,
-   "img_acc": 0.79123,
-   "s1_dE": 0.11822,
-   "s1_dL": 0.11502,
-   "s1_con": 0.25328
+   "step": 8734,
+   "secs": 1257,
+   "loss": 0.12307,
+   "dE": 0.03913,
+   "con": 0.02326,
+   "elem": 0.22178,
+   "e_ok": 0.94556,
+   "s0_dE": 0.03785,
+   "s0_dL": 0.03619,
+   "s0_con": 0.02286,
+   "e_acc": 0.88683,
+   "e_nll": 0.31044,
+   "r_acc": 0.78472,
+   "s1_dE": 0.0992,
+   "s1_dL": 0.09587,
+   "s1_con": 0.19078
   },
   {
    "epoch": 11,
-   "step": 5616,
-   "secs": 822,
-   "loss": 0.14645,
-   "dE": 0.04703,
-   "con": 0.03983,
-   "img": 0.78304,
-   "s0_dE": 0.04345,
-   "s0_dL": 0.04139,
-   "s0_con": 0.01948,
-   "img_acc": 0.81053,
-   "s1_dE": 0.11486,
-   "s1_dL": 0.11153,
-   "s1_con": 0.20122
+   "step": 9528,
+   "secs": 1369,
+   "loss": 0.1209,
+   "dE": 0.03851,
+   "con": 0.02494,
+   "elem": 0.21819,
+   "e_ok": 0.94488,
+   "s0_dE": 0.03751,
+   "s0_dL": 0.03585,
+   "s0_con": 0.01708,
+   "e_acc": 0.88586,
+   "e_nll": 0.31137,
+   "r_acc": 0.79003,
+   "s1_dE": 0.09999,
+   "s1_dL": 0.09671,
+   "s1_con": 0.22857
   },
   {
    "epoch": 12,
-   "step": 6084,
-   "secs": 891,
-   "loss": 0.13803,
-   "dE": 0.04635,
-   "con": 0.03723,
-   "img": 0.76961,
-   "s0_dE": 0.04328,
-   "s0_dL": 0.04125,
-   "s0_con": 0.01831,
-   "img_acc": 0.8051,
-   "s1_dE": 0.11435,
-   "s1_dL": 0.11091,
-   "s1_con": 0.16656
+   "step": 10322,
+   "secs": 1480,
+   "loss": 0.11961,
+   "dE": 0.03832,
+   "con": 0.02343,
+   "elem": 0.21455,
+   "e_ok": 0.94433,
+   "s0_dE": 0.03703,
+   "s0_dL": 0.03539,
+   "s0_con": 0.01577,
+   "e_acc": 0.88654,
+   "e_nll": 0.32437,
+   "r_acc": 0.78799,
+   "s1_dE": 0.10176,
+   "s1_dL": 0.09848,
+   "s1_con": 0.18905
   },
   {
    "epoch": 13,
-   "step": 6552,
-   "secs": 958,
-   "loss": 0.13627,
-   "dE": 0.04589,
-   "con": 0.03889,
-   "img": 0.76454,
-   "s0_dE": 0.04247,
-   "s0_dL": 0.04047,
-   "s0_con": 0.01875,
-   "img_acc": 0.80161,
-   "s1_dE": 0.11794,
-   "s1_dL": 0.11469,
-   "s1_con": 0.23754
+   "step": 11116,
+   "secs": 1602,
+   "loss": 0.11784,
+   "dE": 0.03781,
+   "con": 0.02268,
+   "elem": 0.21099,
+   "e_ok": 0.94429,
+   "s0_dE": 0.03681,
+   "s0_dL": 0.03517,
+   "s0_con": 0.01501,
+   "e_acc": 0.88403,
+   "e_nll": 0.32853,
+   "r_acc": 0.78867,
+   "s1_dE": 0.09691,
+   "s1_dL": 0.09349,
+   "s1_con": 0.1905
   },
   {
    "epoch": 14,
-   "step": 7020,
-   "secs": 1025,
-   "loss": 0.13324,
-   "dE": 0.045,
-   "con": 0.0383,
-   "img": 0.75126,
-   "s0_dE": 0.04194,
-   "s0_dL": 0.03998,
-   "s0_con": 0.02424,
-   "img_acc": 0.80332,
-   "s1_dE": 0.11805,
-   "s1_dL": 0.11474,
-   "s1_con": 0.22873
+   "step": 11910,
+   "secs": 1713,
+   "loss": 0.11638,
+   "dE": 0.03747,
+   "con": 0.02049,
+   "elem": 0.20897,
+   "e_ok": 0.94437,
+   "s0_dE": 0.03681,
+   "s0_dL": 0.03517,
+   "s0_con": 0.01695,
+   "e_acc": 0.88599,
+   "e_nll": 0.32894,
+   "r_acc": 0.78876,
+   "s1_dE": 0.10008,
+   "s1_dL": 0.09666,
+   "s1_con": 0.17253
   },
   {
    "epoch": 15,
-   "step": 7488,
-   "secs": 1093,
-   "loss": 0.13267,
-   "dE": 0.04502,
-   "con": 0.0394,
-   "img": 0.74724,
-   "s0_dE": 0.04183,
-   "s0_dL": 0.03983,
-   "s0_con": 0.01954,
-   "img_acc": 0.81353,
-   "s1_dE": 0.11306,
-   "s1_dL": 0.10954,
-   "s1_con": 0.20409
+   "step": 12704,
+   "secs": 1824,
+   "loss": 0.11504,
+   "dE": 0.03695,
+   "con": 0.02021,
+   "elem": 0.20677,
+   "e_ok": 0.94463,
+   "s0_dE": 0.03671,
+   "s0_dL": 0.03509,
+   "s0_con": 0.01521,
+   "e_acc": 0.88645,
+   "e_nll": 0.33489,
+   "r_acc": 0.78799,
+   "s1_dE": 0.09822,
+   "s1_dL": 0.09477,
+   "s1_con": 0.17304
   },
   {
    "epoch": 16,
-   "step": 7956,
-   "secs": 1161,
-   "loss": 0.1305,
-   "dE": 0.04465,
-   "con": 0.03933,
-   "img": 0.72731,
-   "s0_dE": 0.04195,
-   "s0_dL": 0.04006,
-   "s0_con": 0.02284,
-   "img_acc": 0.80485,
-   "s1_dE": 0.11437,
-   "s1_dL": 0.11122,
-   "s1_con": 0.27215
+   "step": 13498,
+   "secs": 1934,
+   "loss": 0.1141,
+   "dE": 0.03677,
+   "con": 0.02107,
+   "elem": 0.20373,
+   "e_ok": 0.94208,
+   "s0_dE": 0.03658,
+   "s0_dL": 0.03496,
+   "s0_con": 0.01736,
+   "e_acc": 0.88369,
+   "e_nll": 0.33225,
+   "r_acc": 0.78519,
+   "s1_dE": 0.09748,
+   "s1_dL": 0.09415,
+   "s1_con": 0.19973
   },
   {
    "epoch": 17,
-   "step": 8424,
-   "secs": 1229,
-   "loss": 0.12864,
-   "dE": 0.04415,
-   "con": 0.03796,
-   "img": 0.71419,
-   "s0_dE": 0.04177,
-   "s0_dL": 0.0399,
-   "s0_con": 0.0207,
-   "img_acc": 0.79611,
-   "s1_dE": 0.10924,
-   "s1_dL": 0.10592,
-   "s1_con": 0.22063
-  },
-  {
-   "epoch": 18,
-   "step": 8892,
-   "secs": 1296,
-   "loss": 0.12691,
-   "dE": 0.04379,
-   "con": 0.03646,
-   "img": 0.69736,
-   "s0_dE": 0.04175,
-   "s0_dL": 0.03993,
-   "s0_con": 0.02416,
-   "img_acc": 0.801,
-   "s1_dE": 0.10774,
-   "s1_dL": 0.10449,
-   "s1_con": 0.24435
-  },
-  {
-   "epoch": 19,
-   "step": 9360,
-   "secs": 1364,
-   "loss": 0.12499,
-   "dE": 0.04319,
-   "con": 0.036,
-   "img": 0.68816,
-   "s0_dE": 0.0413,
-   "s0_dL": 0.03945,
-   "s0_con": 0.02047,
-   "img_acc": 0.79813,
-   "s1_dE": 0.11372,
-   "s1_dL": 0.11037,
-   "s1_con": 0.20421
-  },
-  {
-   "epoch": 20,
-   "step": 9828,
-   "secs": 1432,
-   "loss": 0.12434,
-   "dE": 0.04334,
-   "con": 0.03581,
-   "img": 0.67546,
-   "s0_dE": 0.0409,
-   "s0_dL": 0.03905,
-   "s0_con": 0.01907,
-   "img_acc": 0.79483,
-   "s1_dE": 0.10692,
-   "s1_dL": 0.1036,
-   "s1_con": 0.22495
-  },
-  {
-   "epoch": 21,
-   "step": 10296,
-   "secs": 1499,
-   "loss": 0.12248,
-   "dE": 0.04287,
-   "con": 0.03581,
-   "img": 0.65995,
-   "s0_dE": 0.04105,
-   "s0_dL": 0.0392,
-   "s0_con": 0.02163,
-   "img_acc": 0.7999,
-   "s1_dE": 0.10461,
-   "s1_dL": 0.10127,
-   "s1_con": 0.18494
-  },
-  {
-   "epoch": 22,
-   "step": 10764,
-   "secs": 1567,
-   "loss": 0.12033,
-   "dE": 0.04222,
-   "con": 0.03395,
-   "img": 0.64253,
-   "s0_dE": 0.04075,
-   "s0_dL": 0.03893,
-   "s0_con": 0.02266,
-   "img_acc": 0.7999,
-   "s1_dE": 0.10178,
-   "s1_dL": 0.09855,
-   "s1_con": 0.23005
-  },
-  {
-   "epoch": 23,
-   "step": 11232,
-   "secs": 1635,
-   "loss": 0.11905,
-   "dE": 0.04192,
-   "con": 0.03691,
-   "img": 0.63201,
-   "s0_dE": 0.04024,
-   "s0_dL": 0.03843,
-   "s0_con": 0.01664,
-   "img_acc": 0.79049,
-   "s1_dE": 0.10458,
-   "s1_dL": 0.10125,
-   "s1_con": 0.1973
-  },
-  {
-   "epoch": 24,
-   "step": 11700,
-   "secs": 1703,
-   "loss": 0.11767,
-   "dE": 0.04169,
-   "con": 0.03462,
-   "img": 0.61761,
-   "s0_dE": 0.04043,
-   "s0_dL": 0.03861,
-   "s0_con": 0.02038,
-   "img_acc": 0.79379,
-   "s1_dE": 0.1025,
-   "s1_dL": 0.09926,
-   "s1_con": 0.2135
-  },
-  {
-   "epoch": 25,
-   "step": 12168,
-   "secs": 1770,
-   "loss": 0.11655,
-   "dE": 0.04131,
-   "con": 0.03564,
-   "img": 0.60858,
-   "s0_dE": 0.04068,
-   "s0_dL": 0.0389,
-   "s0_con": 0.01981,
-   "img_acc": 0.79715,
-   "s1_dE": 0.10372,
-   "s1_dL": 0.10038,
-   "s1_con": 0.20402
-  },
-  {
-   "epoch": 26,
-   "step": 12636,
-   "secs": 1838,
-   "loss": 0.11502,
-   "dE": 0.04106,
-   "con": 0.03415,
-   "img": 0.59825,
-   "s0_dE": 0.04032,
-   "s0_dL": 0.03852,
-   "s0_con": 0.01906,
-   "img_acc": 0.79691,
-   "s1_dE": 0.10293,
-   "s1_dL": 0.09969,
-   "s1_con": 0.22009
-  },
-  {
-   "epoch": 27,
-   "step": 13104,
-   "secs": 1906,
-   "loss": 0.11521,
-   "dE": 0.04109,
-   "con": 0.03603,
-   "img": 0.59315,
-   "s0_dE": 0.0404,
-   "s0_dL": 0.03861,
-   "s0_con": 0.02053,
-   "img_acc": 0.79721,
-   "s1_dE": 0.10302,
-   "s1_dL": 0.09978,
-   "s1_con": 0.2078
-  },
-  {
-   "epoch": 28,
-   "step": 13572,
-   "secs": 1973,
-   "loss": 0.11402,
-   "dE": 0.04082,
-   "con": 0.03683,
-   "img": 0.58722,
-   "s0_dE": 0.04034,
-   "s0_dL": 0.03856,
-   "s0_con": 0.01831,
-   "img_acc": 0.79098,
-   "s1_dE": 0.10383,
-   "s1_dL": 0.10053,
-   "s1_con": 0.181
-  },
-  {
-   "epoch": 29,
-   "step": 14040,
-   "secs": 2041,
-   "loss": 0.11321,
-   "dE": 0.04059,
-   "con": 0.03495,
-   "img": 0.5803,
-   "s0_dE": 0.04018,
-   "s0_dL": 0.03841,
-   "s0_con": 0.01978,
-   "img_acc": 0.79416,
-   "s1_dE": 0.10315,
-   "s1_dL": 0.09986,
-   "s1_con": 0.19971
+   "step": 14292,
+   "secs": 2053,
+   "loss": 0.11284,
+   "dE": 0.03617,
+   "con": 0.02163,
+   "elem": 0.20522,
+   "e_ok": 0.94378,
+   "s0_dE": 0.03654,
+   "s0_dL": 0.03492,
+   "s0_con": 0.01713,
+   "e_acc": 0.88535,
+   "e_nll": 0.34029,
+   "r_acc": 0.78514,
+   "s1_dE": 0.09664,
+   "s1_dL": 0.09326,
+   "s1_con": 0.20492
   }
  ],
- "best_score": 0.07564117624840601
+ "best_score": 0.0686828651548906
 }

runtime/darkreader-nadir.js

@@ -3229,6 +3229,10 @@
               }
             : null;
         return {
+            // Nadir: dark vs near-white pixels (a logo's white halo or plate is not ink; colours and anti-aliased
+            // edges are neither), and the dark share of all opaque pixels
+            inkDarkShare: darkPixelsCount / Math.max(1, darkPixelsCount + lightPixelsCount),
+            darkShare: darkPixelsCount / Math.max(1, opaquePixelsCount),
             isDark: darkPixelsCount / opaquePixelsCount >= DARK_IMAGE_THRESHOLD,
             isLight:
                 lightPixelsCount / opaquePixelsCount >= LIGHT_IMAGE_THRESHOLD,
@@ -3973,7 +3977,7 @@
                         }
                     }
                     if (imageDetails) {
-                        const bgImageValue = getBgImageValue(
+                        const bgImageValue = await getBgImageValue(
                             imageDetails,
                             theme
                         );
@@ -4005,7 +4009,7 @@
                     size.includes("100%");
                 return !isTiled && !isStretched;
             };
-            const getBgImageValue = (imageDetails, theme) => {
+            const getBgImageValue = async (imageDetails, theme) => {
                 const {
                     isDark,
                     isLight,
@@ -4022,7 +4026,9 @@
                 if (nadirImageHook && theme.mode === 1) {
                     let decision = null;
                     try {
-                        decision = nadirImageHook(imageDetails.src, {
+                        // Nadir: the element head decides; for an image it has not seen yet the hook may return a
+                        // promise (measure + classify), or an adapted copy as {url}
+                        decision = await nadirImageHook(imageDetails.src, {
                             isDark,
                             isLight,
                             isTransparent,
@@ -4035,7 +4041,8 @@
                             ).toLowerCase(),
                             repeat: (
                                 rule.style.backgroundRepeat || ""
-                            ).toLowerCase()
+                            ).toLowerCase(),
+                            selector: rule.selectorText || ""
                         });
                     } catch (err) {
                         decision = null;
@@ -4043,6 +4050,14 @@
                     if (decision === "none") {
                         return "none";
                     }
+                    if (decision && typeof decision === "object" && decision.url) {
+                        return `url("${decision.url}")`;
+                    }
+                    if (decision === "keep") {
+                        // the element head chose to show it as it is (the solid-colour test below misreads
+                        // SVGs without an intrinsic size, e.g. a blue tile with a white glyph)
+                        return null;
+                    }
                     if (solidColor) {
                         const darkColor = modifyBackgroundColor(
                             solidColor,
@@ -6499,8 +6514,14 @@
                     );
                     const dataURL = `data:image/svg+xml;base64,${btoa(svgString)}`;
                     getImageDetails(dataURL).then((details) => {
+                        // Nadir: dark letters with a white halo or on a white plate are a dark logo too
+                        const darkInk =
+                            details.isDark ||
+                            (globalThis.__nadirFillSide &&
+                                details.inkDarkShare >= 0.55 &&
+                                details.darkShare >= 0.3);
                         if (
-                            (details.isDark && details.isTransparent) ||
+                            (darkInk && details.isTransparent) ||
                             (details.isLarge &&
                                 details.isLight &&
                                 !details.isTransparent)

runtime/extract.js

@@ -228,9 +228,25 @@
     return t;
   }
 
+  // The part of a CSS background image that is shown, in CSS px: [x, y, w, h, displayed image w, h] (0 = natural size).
+  // Only plain px positions/sizes (icon sprites); keyword/percentage layouts show the whole image.
+  function bgCrop(cs, w, h) {
+    const pos = (cs.backgroundPosition || '').split(',')[0].trim().split(/\s+/);
+    const size = (cs.backgroundSize || 'auto').split(',')[0].trim().split(/\s+/);
+    if (pos.length !== 2 || !pos.every(v => /^-?[\d.]+px$/.test(v))) return null;
+    const px = parseFloat(pos[0]), py = parseFloat(pos[1]);
+    let sw = 0, sh = 0;
+    if (size.length === 2 && size.every(v => /^[\d.]+px$/.test(v))) { sw = parseFloat(size[0]); sh = parseFloat(size[1]); }
+    else if (size.length === 1 && /^[\d.]+px$/.test(size[0])) { sw = parseFloat(size[0]); sh = -1; } // height follows the aspect ratio
+    else if (!(size.length === 1 && size[0] === 'auto') && !(size.length === 2 && size[0] === 'auto' && size[1] === 'auto')) return null;
+    if (px >= 0 && py >= 0) return null; // drawn from its corner: the whole image (or a texture) is shown
+    return [Math.round(-px), Math.round(-py), Math.round(w), Math.round(h), sw, sh];
+  }
+
   async function extract(opts) {
     opts = Object.assign({maxElements: 5000, maxRules: 40000, maxTokens: 384, maxImages: 48, describe: false,
-      thumbs: false, thumbSize: 32, bigThumbSize: 0, fetchImage: null, pseudo: true, time: true}, opts || {});
+      thumbs: false, thumbSize: 32, bigThumbSize: 0, fetchImage: null, pseudo: true, time: true,
+      srcMax: 600}, opts || {}); // crawl records keep 600 chars of a URL; the extension needs whole data: URLs to decode them
     const t0 = performance.now();
     const vw = innerWidth, vh = innerHeight;
     const sx = scrollX, sy = scrollY;
@@ -327,6 +343,23 @@
     const rootRec = {surf: rootSurface, surfTok: rootTok, imgBg: false, region: 0, depth: 0, fixed: false};
     const images = [];
     const svgRoots = [];
+    // an image drawn with `content: url(...)` in ::before/::after (icon fonts' image cousins, sprite icons)
+    function pseudoImage(el, parent, pseudo, ct, rect, x, y, area1, rec, vis) {
+      const pu = ct && /^url\(\s*["']?([^"')]+)["']?\s*\)/.exec(ct);
+      if (!pu) return;
+      let bx = x, by = y, bw = rect.width, bh = rect.height, cropPx = null;
+      const pe = parent && !(parent instanceof SVGElement) ? parent : null;
+      if ((bw > 72 || bh > 72) && pe) { // one icon of a sprite sheet: the element is shifted inside a small clipping box
+        const prc = pe.getBoundingClientRect();
+        if (prc.width <= 72 && prc.height <= 72 && prc.width >= 4 && prc.height >= 4 && getComputedStyle(pe).overflow !== 'visible') {
+          cropPx = [Math.round(prc.left + sx - x), Math.round(prc.top + sy - y), Math.round(prc.width), Math.round(prc.height), 0, 0];
+          bx = prc.left + sx; by = prc.top + sy; bw = prc.width; bh = prc.height;
+        }
+      }
+      if (!cropPx && !vis) return;
+      if (bw * bh <= 16 || bw > 400 || bh > 400) return;
+      try { images.push({el, kind: 'pseudo', src: new URL(pu[1], location.href).href, x: bx, y: by, w: bw, h: bh, area: bw * bh, area1: Math.min(area1 || bw * bh, bw * bh), rec, cropPx, pseudo}); } catch (e) {}
+    }
     let visibleCount = 0;
 
     const ec = opts.elemColors ? [] : null;
@@ -411,7 +444,7 @@
         }
         const urls = [...bgImg.matchAll(/url\(\s*["']?([^"')]+)["']?\s*\)/g)].map(m => m[1]);
         if (urls.length && vis && area > 64) {
-          images.push({el, kind: 'bg', src: urls[0], x, y, w, h, area, area1, rec});
+          images.push({el, kind: 'bg', src: urls[0], x, y, w, h, area, area1, rec, cropPx: bgCrop(cs, w, h)});
           if (area > 0.02 * vw * vh) rec.imgBg = true;
         }
       }
@@ -509,6 +542,10 @@
         const href = el.getAttribute('href') || el.getAttribute('xlink:href');
         if (href) images.push({el, kind: 'svgimage', src: new URL(href, location.href).href, x, y, w, h, area, area1, rec});
       }
+      // sprite-sheet icons whose (huge, shifted) element box counts as off-page: only the clipped part shows
+      if (opts.pseudo && !vis && !hidden && !isSvg && (rect.width > 72 || rect.height > 72)) {
+        for (const pseudo of ['::before', '::after']) pseudoImage(el, parent, pseudo, getComputedStyle(el, pseudo).content, rect, x, y, area1, rec, false);
+      }
       // pseudo elements
       if (opts.pseudo && vis && !isSvg && els.length < 20000) {
         for (const pseudo of ['::before', '::after']) {
@@ -518,6 +555,7 @@
           const pb = parseColor(ps.backgroundColor);
           const pw = parseFloat(ps.width) || 0, ph = parseFloat(ps.height) || 0;
           const parea = Math.min(area, pw * ph);
+          pseudoImage(el, parent, pseudo, ct, rect, x, y, area1, rec, true);
           if (pb) { const t = tok(R.bg, pb); flags(t); t.area += parea; addRel('bb', t, rec.surfTok, parea); }
           if (ct !== '""' && ct !== "''") {
             const pc = parseColor(ps.color);
@@ -637,7 +675,7 @@
       if (imgOut.length >= opts.maxImages) break;
       if (im.kind !== 'svg' && (!im.src || im.src.startsWith('data:image/gif;base64,R0lGOD'))) continue;
       if (im.kind === 'svg') { im.src = 'svg:' + svgCount++; if (opts.markSvg) im.el.setAttribute('data-nadir-svg', im.src); }
-      const key = im.src;
+      const key = im.src + (im.cropPx ? '#nadir-crop=' + im.cropPx.join(',') : '');
       if (imgSeen.has(key)) { const o = imgOut[imgSeen.get(key)]; o.count++; o.areaSum += im.area; continue; }
       const el = im.el;
       const a = im.rec.link;
@@ -647,12 +685,13 @@
       const meta = im.kind === 'svg' ? (im.svgMeta || '') : ((el.getAttribute && (el.getAttribute('alt') || '')) + ' ' + (typeof el.className === 'string' ? el.className : '') + ' ' + (el.id || '') + ' ' + (im.src.split('?')[0].split('/').pop() || '')).toLowerCase();
       const ext = im.kind === 'svg' ? 'svg' : (im.src.split('?')[0].match(/\.(svg|png|jpe?g|webp|gif|avif|ico)$/i) || [, ''])[1].toLowerCase();
       const under = oklab(im.rec.surf);
-      const o = {src: im.src.slice(0, 600), kind: im.kind, x: Math.round(im.x), y: Math.round(im.y), w: Math.round(im.w), h: Math.round(im.h), count: 1, areaSum: im.area,
+      const o = {src: im.src.slice(0, opts.srcMax), kind: im.kind, x: Math.round(im.x), y: Math.round(im.y), w: Math.round(im.w), h: Math.round(im.h), count: 1, areaSum: im.area,
         nw: el.naturalWidth || 0, nh: el.naturalHeight || 0, homeLink, inLink: im.rec.inLink ? 1 : 0, header: (im.rec.region & 3) ? 1 : 0, footer: (im.rec.region & 4) ? 1 : 0,
         logo: /logo|brand|wordmark/.test(meta) ? 1 : 0, icon: /icon|ico\b|glyph|arrow|chevron|caret|svg/.test(meta) ? 1 : 0, avatar: /avatar|profile|user/.test(meta) ? 1 : 0,
         ext, underTok: im.rec.surfTok && im.rec.surfTok.i >= 0 ? im.rec.surfTok.i : -1, under,
         surf: im.rec.surf.slice(0, 3).map(Math.round), vy: Math.round(im.y - sy), vx: Math.round(im.x - sx), useExt: im.useExt ? 1 : 0,
         alt: opts.describe ? ((el.getAttribute && el.getAttribute('alt')) || '').slice(0, 60) : undefined,
+        cropPx: im.cropPx || undefined, pseudo: im.pseudo || undefined,
         stats: null, thumb: null};
       if (im.kind === 'img' && el.complete === false) o.loading = 1;
       imgSeen.set(key, imgOut.length);
@@ -661,7 +700,7 @@
     if (opts.thumbs) await imageStats(imgOut, opts);
     for (const o of imgOut) o.f = imageFeatures(o, vw, vh, docH);
     const res = {
-      v: 5, url: location.href, title: document.title.slice(0, 200), vw, vh, docH, docW,
+      v: 6, url: location.href, title: document.title.slice(0, 200), vw, vh, docH, docW,
       scheme: schemeInfo(sheetScan), darkness, canvasDark,
       tokens: out, rel: relOut, images: imgOut,
       ec: ec || undefined,
@@ -714,9 +753,21 @@
       try {
         const img = await loadImage(o.src, opts.fetchImage);
         if (!img) return;
-        const IW = img.naturalWidth || img.width, IH = img.naturalHeight || img.height;
+        let IW = img.naturalWidth || img.width, IH = img.naturalHeight || img.height, src = img;
         o.nw = o.nw || IW; o.nh = o.nh || IH;
-        statsFromSource(o, img, IW, IH, opts);
+        let crop = null;
+        if (o.cropPx) { // one icon of a sprite sheet: measure only the part that is shown
+          const [cx, cy, cw, ch, sw, sh] = o.cropPx;
+          const kx = sw > 0 ? IW / sw : 1, ky = sh > 0 ? IH / sh : (sh < 0 ? kx : 1);
+          crop = [cx * kx, cy * ky, cw * kx, ch * ky];
+        } else if (o.cropFrac) { const [fx, fy, fw, fh] = o.cropFrac; crop = [fx * IW, fy * IH, fw * IW, fh * IH]; }
+        if (crop) {
+          const c = document.createElement('canvas');
+          c.width = Math.max(1, Math.round(crop[2])); c.height = Math.max(1, Math.round(crop[3]));
+          c.getContext('2d').drawImage(img, crop[0], crop[1], crop[2], crop[3], 0, 0, c.width, c.height);
+          src = c; IW = c.width; IH = c.height;
+        }
+        statsFromSource(o, src, IW, IH, opts);
         if (img.close) img.close();
       } catch (e) { o.err = String(e).slice(0, 80); }
     });
@@ -769,6 +820,7 @@
       o.under ? o.under[0] : 1, o.under ? o.under[1] * 4 : 0, o.under ? o.under[2] * 4 : 0,
       o.stats ? 1 : 0, s.transp || 0, s.semi || 0, s.meanL || 0, s.stdL || 0, (s.meanC || 0) * 4, s.white || 0, s.black || 0, s.ringA || 0, s.ringL || 0, s.colors || 0,
       o.kind === 'svg' ? 1 : 0, o.useExt || 0,
+      o.kind === 'pseudo' ? 1 : 0, o.cropPx ? 1 : 0,
     ]);
   }
   const sanitize = (f) => f.map((v) => (Number.isFinite(+v) ? +v : 0));
@@ -817,6 +869,59 @@
     statsFromSource(o, c, c.width, c.height, opts);
   }
 
+  // ---------- one image found after the page was analysed (late-loaded, CSS background, sprite or pseudo icon) ----------
+  // Same record and features as extract() builds in its walk: the surface under it (backgrounds and gradients of the
+  // element and its ancestors composited over the canvas), link/region flags, logo/icon/avatar naming.
+  function surfaceOf(el) {
+    const chain = [];
+    for (let e = el; e && e.nodeType === 1; e = e.parentElement || (e.getRootNode && e.getRootNode().host) || null) chain.push(e);
+    const rootCS = getComputedStyle(document.documentElement);
+    const canvasDark = /dark/.test(rootCS.colorScheme) && !/light/.test(rootCS.colorScheme);
+    let surf = canvasDark ? [18, 18, 18, 1] : [255, 255, 255, 1];
+    for (let i = chain.length - 1; i >= 0; i--) {
+      const cs = getComputedStyle(chain[i]);
+      const bg = parseColor(cs.backgroundColor);
+      if (bg) surf = over(bg, surf);
+      const bi = cs.backgroundImage;
+      if (bi && bi !== 'none' && bi.includes('gradient')) {
+        const gc = colorsIn(bi);
+        if (gc.length) {
+          const avg = [0, 0, 0, 1];
+          for (const c of gc) { const o = over(c, surf); avg[0] += o[0] / gc.length; avg[1] += o[1] / gc.length; avg[2] += o[2] / gc.length; }
+          surf = avg;
+        }
+      }
+    }
+    return surf;
+  }
+  function recordFor(el, kind, src, cropPx, box) {
+    const vw = innerWidth, vh = innerHeight, sx = scrollX, sy = scrollY;
+    const docEl = document.documentElement;
+    const docH = Math.max(docEl.scrollHeight, document.body ? document.body.scrollHeight : 0, vh);
+    const r = box || el.getBoundingClientRect();
+    const x = r.left + sx, y = r.top + sy, w = r.width, h = r.height;
+    let region = 0;
+    for (let e = el; e && e.nodeType === 1; e = e.parentElement) {
+      const role = e.getAttribute && e.getAttribute('role');
+      if (e.tagName === 'HEADER' || role === 'banner') region |= 1;
+      if (e.tagName === 'NAV' || role === 'navigation') region |= 2;
+      if (e.tagName === 'FOOTER' || role === 'contentinfo') region |= 4;
+    }
+    const a = el.closest ? el.closest('a') : null;
+    let homeLink = 0;
+    try { if (a) { const u = new URL(a.getAttribute('href') || '', location.href); homeLink = (u.origin === location.origin && (u.pathname === '/' || u.pathname === '')) ? 1 : 0; } } catch (e) {}
+    const meta = (((el.getAttribute && el.getAttribute('alt')) || '') + ' ' + (typeof el.className === 'string' ? el.className : '') + ' ' + (el.id || '') + ' ' + (src.split('?')[0].split('/').pop() || '')).toLowerCase();
+    const ext = (src.split('?')[0].match(/\.(svg|png|jpe?g|webp|gif|avif|ico)$/i) || [, ''])[1].toLowerCase();
+    const surf = surfaceOf(el);
+    return {src, kind, x: Math.round(x), y: Math.round(y), w: Math.round(w), h: Math.round(h), count: 1, areaSum: w * h,
+      nw: el.naturalWidth || 0, nh: el.naturalHeight || 0, homeLink, inLink: a ? 1 : 0, header: (region & 3) ? 1 : 0, footer: (region & 4) ? 1 : 0,
+      logo: /logo|brand|wordmark/.test(meta) ? 1 : 0, icon: /icon|ico\b|glyph|arrow|chevron|caret|svg/.test(meta) ? 1 : 0, avatar: /avatar|profile|user/.test(meta) ? 1 : 0,
+      ext, underTok: -1, under: oklab(surf), surf: surf.slice(0, 3).map(Math.round), vy: Math.round(y - sy), vx: Math.round(x - sx), useExt: 0,
+      cropPx: cropPx || undefined, stats: null, thumb: null, dims: {vw, vh, docH}};
+  }
+
+  globalThis.__nadirImageRecord = recordFor;
+  globalThis.__nadirBgCrop = bgCrop;
   globalThis.__nadirImageStats = imageStatsStandalone;
   globalThis.__nadirImageFeatures = imageFeatures;
   globalThis.__nadirExtract = extract;

runtime/nadir-infer.js

@@ -2,8 +2,10 @@
 // Mirrors train.py collate() for a single page. Works in browsers (offscreen document) and Node (parity tests).
 (function (root) {
   'use strict';
-  const MAX_TOK = 256, MAX_IMG = 32, FT = 49, FI = 39, K = 11;
-  const TREAT = ['keep', 'dim', 'invert', 'light', 'plate'];
+  const MAX_TOK = 256, MAX_IMG = 32, FT = 49, FI = 41, K = 11, N_CTX = 32;
+  const TREAT = ['keep', 'invert', 'adapt', 'remove'];
+  const IMG_ROLES = ['photo', 'illustration', 'diagram', 'logo', 'icon', 'text', 'backdrop', 'avatar', 'ui'];
+  const padF = (f) => { const o = new Array(FI).fill(0); for (let i = 0; i < Math.min(FI, f.length); i++) o[i] = Number.isFinite(f[i]) ? f[i] : 0; return o; };
   const FG_ROLES = new Set([1, 5, 6]); // text, svg fill, svg stroke
   const lin = (u) => { u /= 255; return u <= 0.04045 ? u / 12.92 : Math.pow((u + 0.055) / 1.055, 2.4); };
   function oklab(r, g, b) {
@@ -21,8 +23,56 @@
   }
   const relLum = (rgb) => 0.2126 * lin(rgb[0]) + 0.7152 * lin(rgb[1]) + 0.0722 * lin(rgb[2]);
   const contrast = (a, b) => { const x = relLum(a), y = relLum(b); return (Math.max(x, y) + 0.05) / (Math.min(x, y) + 0.05); };
+  // WCAG target for a foreground token: 3:1 for large/bold text and graphics, 4.5:1 for body text
+  const needOf = (t) => {
+    if (t.role !== 1) return 3;
+    const fs = t.f[21] * 32, bold = t.f[23] > 0.5;
+    return (fs >= 24 || (fs >= 18.5 && bold)) ? 3 : 4.5;
+  };
   const hex = (rgb, a) => '#' + rgb.map(v => v.toString(16).padStart(2, '0')).join('') + (a < 0.999 ? Math.round(a * 255).toString(16).padStart(2, '0') : '');
 
+  function putThumb(dst, j, thumb) { // RGBA bytes (32x32, row-major) -> planar floats
+    if (!thumb || thumb.length !== 4096) return;
+    for (let y = 0; y < 32; y++) for (let x = 0; x < 32; x++) for (let c = 0; c < 4; c++)
+      dst[j * 4096 + c * 1024 + y * 32 + x] = thumb[(y * 32 + x) * 4 + c] / 255;
+  }
+
+  // element head inputs: items [{f, thumb, ul: OKLab under (light), ud: OKLab under (dark)}], ctx Float32Array(32)
+  function buildElemInputs(items, ctx) {
+    const Kn = Math.max(1, items.length);
+    const e_f = new Float32Array(Kn * FI), e_thumb = new Float32Array(Kn * 4096), ul = new Float32Array(Kn * 3), ud = new Float32Array(Kn * 3);
+    items.forEach((it, k) => { e_f.set(padF(it.f), k * FI); putThumb(e_thumb, k, it.thumb); ul.set(it.ul, k * 3); ud.set(it.ud, k * 3); });
+    return {K: items.length, feeds: {e_f: [e_f, [1, Kn, FI]], e_thumb: [e_thumb, [1, Kn, 4, 32, 32]], under_light: [ul, [1, Kn, 3]],
+      under_dark: [ud, [1, Kn, 3]], ctx: [Float32Array.from(ctx), [1, N_CTX]]}};
+  }
+  // calibrated decisions: [{t: treatment, p: its probability, role, rp: role probability}]
+  function decodeElements(treatLogits, roleLogits, K, temperature) {
+    const T = temperature || 1, out = [];
+    for (let k = 0; k < K; k++) {
+      const z = Array.from(treatLogits.slice(k * 4, k * 4 + 4), (v) => v / T);
+      const mx = Math.max(...z), e = z.map((v) => Math.exp(v - mx)), s = e.reduce((a, b) => a + b, 0);
+      let best = 0; for (let c = 1; c < 4; c++) if (e[c] > e[best]) best = c;
+      const r = Array.from(roleLogits.slice(k * 9, k * 9 + 9));
+      const rm = Math.max(...r), re = r.map((v) => Math.exp(v - rm)), rs = re.reduce((a, b) => a + b, 0);
+      let rb = 0; for (let c = 1; c < 9; c++) if (re[c] > re[rb]) rb = c;
+      out.push({t: TREAT[best], p: +(e[best] / s).toFixed(3), probs: e.map((v) => +(v / s).toFixed(3)), role: IMG_ROLES[rb], rp: +(re[rb] / rs).toFixed(3)});
+    }
+    return out;
+  }
+  // the page's own images: the dark surface under each one is the palette head's raw output for its surface token
+  function pageElementItems(page, out) {
+    const toks = page.tokens.slice(0, MAX_TOK), n = toks.length;
+    let root = 0; toks.forEach((t, i) => { if (t.f[24] > toks[root].f[24]) root = i; });
+    const items = [], idx = [];
+    page.images.forEach((im, j) => {
+      if (!im.thumb || im.thumb.length !== 4096 || !im.f) return;
+      const u = im.underTok != null && im.underTok >= 0 && im.underTok < n ? im.underTok : root;
+      items.push({f: im.f, thumb: im.thumb, ul: [im.f[23], im.f[24] / 4, im.f[25] / 4], ud: Array.from(out.lab.slice(u * 3, u * 3 + 3)), u});
+      idx.push(j);
+    });
+    return {items, idx};
+  }
+
   function buildInputs(page) {
     const toks = page.tokens.slice(0, MAX_TOK);
     const N = Math.max(1, toks.length);
@@ -36,11 +86,8 @@
     });
     const img_f = new Float32Array(M * FI), thumbs = new Float32Array(M * 4 * 32 * 32);
     imgs.forEach((im, j) => {
-      img_f.set(im.f.slice(0, FI).map((v) => (Number.isFinite(v) ? v : 0)), j * FI);
-      if (im.thumb && im.thumb.length === 4096) {
-        for (let y = 0; y < 32; y++) for (let x = 0; x < 32; x++) for (let c = 0; c < 4; c++)
-          thumbs[j * 4096 + c * 1024 + y * 32 + x] = im.thumb[(y * 32 + x) * 4 + c] / 255;
-      }
+      img_f.set(padF(im.f), j * FI);
+      putThumb(thumbs, j, im.thumb);
     });
     const st = page.stats || {};
     const page_f = new Float32Array([page.vw / 1280, Math.min(8, page.docH / page.vh) / 8, toks.length / MAX_TOK, Mreal / MAX_IMG,
@@ -85,7 +132,7 @@
     }};
   }
 
-  // outputs: lab Float32Array(N*3), alpha_logit Float32Array(N), img_logits Float32Array(M*5)
+  // outputs: lab Float32Array(N*3), alpha_logit Float32Array(N), ctx Float32Array(32)
   // brand colours keep their hue: a coloured input whose output drifts more than 45 degrees is turned back to its
   // own hue (keeping the predicted lightness and chroma); neutral inputs get at most a faint tint
   function hueGuard(t, o) {
@@ -122,10 +169,28 @@
         outputs[i] = {rgb, a: t.rgba[3], L: oklab(...rgb)[0]};
       } else hueGuard(t, outputs[i]);
     });
+    // a see-through light panel (frosted card, glass button: white at 40-99% opacity) stays see-through but turns
+    // dark; left mid-grey it is a surface no text colour reads on (the contrast pass below only sees opaque ones)
+    toks.forEach((t, i) => {
+      if (keep.has(i) || t.role !== 0 || t.rgba[3] >= 0.999 || t.rgba[3] < 0.4) return;
+      if (oklab(t.rgba[0], t.rgba[1], t.rgba[2])[0] < 0.85 || outputs[i].L <= 0.45) return;
+      const o = oklab(...outputs[i].rgb);
+      outputs[i].rgb = labToRgb(0.24, o[1] * 0.5, o[2] * 0.5); outputs[i].L = 0.24;
+    });
     // safety: page canvas must be dark
     toks.forEach((t, i) => {
       if (t.f[24] > 0.5 && t.role === 0 && outputs[i].L > 0.35) { const o = oklab(...outputs[i].rgb); outputs[i].rgb = labToRgb(0.2, o[1] * 0.3, o[2] * 0.3); outputs[i].L = 0.2; }
     });
+    // an opaque surface stays opaque: see-through, it would show whatever the design stacks beneath it (a hidden
+    // sticky header, the next slide). The predicted colour is composited over the dark canvas instead
+    const canvasI = toks.findIndex((t) => t.f[24] > 0.5 && t.role === 0);
+    const under = canvasI >= 0 ? outputs[canvasI].rgb : [18, 18, 20];
+    toks.forEach((t, i) => {
+      if (keep.has(i) || t.role !== 0 || t.rgba[3] < 0.999 || outputs[i].a >= 0.999) return;
+      const a = outputs[i].a;
+      outputs[i].rgb = outputs[i].rgb.map((c, k) => Math.round(c * a + under[k] * (1 - a)));
+      outputs[i].a = 1; outputs[i].L = oklab(...outputs[i].rgb)[0];
+    });
     // things that were invisible against their surface stay invisible: gradient stops with exactly a surface's
     // colour (text fades) take that surface's dark colour, and so do borders that matched their surface
     const keyOf = (t) => t.rgba[0] * 65536 + t.rgba[1] * 256 + t.rgba[2];
@@ -161,13 +226,12 @@
     for (const [typ, a, b, w] of page.rel) {
       if (typ !== 0 || a >= n || b >= n || toks[b].role !== 0 || outputs[b].a < 0.999) continue;
       if (!surf.has(a)) surf.set(a, []);
-      surf.get(a).push([b, w + 0.1]);
+      surf.get(a).push([b, Math.expm1(w) + 1]); // rels store log1p(weight)
     }
     for (const [i, lst] of surf) {
       const t = toks[i];
       if (keep.has(i) || !FG_ROLES.has(t.role) || outputs[i].a < 0.999) continue;
-      let need0 = 3;
-      if (t.role === 1) { const fs = t.f[21] * 32, bold = t.f[23] > 0.5; need0 = (fs >= 24 || (fs >= 18.5 && bold)) ? 3 : 4.5; }
+      let need0 = needOf(t);
       if (opts.minContrast) need0 = Math.min(need0, opts.minContrast);
       const reqs = [];
       for (const [j, w] of lst) {
@@ -182,24 +246,45 @@
       let best = null;
       for (let k = 0; k <= 50; k++) {
         const Lx = k / 50, rgb = labToRgb(Lx, o[1], o[2]);
-        const c = deficit(rgb) / W + 0.05 * Math.abs(Lx - o[0]);
+        // moving away from the model's lightness costs: the surfaces a colour mostly sits on decide, and a few
+        // bright accent surfaces cannot drag it to a mid-grey that reads nowhere (they are rescued below)
+        const c = deficit(rgb) / W + Math.abs(Lx - o[0]);
         if (!best || c < best[0]) best = [c, rgb, Lx];
       }
       outputs[i].rgb = best[1]; outputs[i].L = best[2];
     }
+    // a light surface whose text no longer reads (its text colour follows the dark page) is darkened within its
+    // own hue until that text reads again: a yellow note becomes dark gold, not a glare
+    const fgOn = new Map();
+    for (const [typ, a, b, w] of page.rel) {
+      if (typ !== 0 || a >= n || b >= n || toks[b].role !== 0 || !FG_ROLES.has(toks[a].role)) continue;
+      if (!fgOn.has(b)) fgOn.set(b, []);
+      fgOn.get(b).push([a, Math.expm1(w) + 1]);
+    }
+    for (const [j, lst] of fgOn) {
+      const s = toks[j];
+      if (keep.has(j) || s.f[24] > 0.5 || outputs[j].a < 0.999 || outputs[j].L < 0.6) continue;
+      const reqs = [];
+      for (const [i, w] of lst) {
+        if (keep.has(i) || outputs[i].a < 0.999) continue;
+        const need = Math.min(needOf(toks[i]), contrast(toks[i].rgba, s.rgba));
+        if (need >= 1.5) reqs.push([outputs[i].rgb, need, w]);
+      }
+      if (!reqs.length) continue;
+      let W = 0; for (const r of reqs) W += r[2];
+      const fail = (rgb) => { let f = 0; for (const [fg, need, w] of reqs) if (contrast(fg, rgb) < need) f += w; return f / W; };
+      if (fail(outputs[j].rgb) <= 0.5) continue;
+      const o = oklab(...outputs[j].rgb);
+      for (let k = Math.floor(o[0] * 50); k >= 10; k--) {
+        const rgb = labToRgb(k / 50, o[1], o[2]);
+        if (fail(rgb) <= 0.2) { outputs[j].rgb = rgb; outputs[j].L = k / 50; break; }
+      }
+    }
     const colors = toks.map((t, i) => keep.has(i) ? null : hex(outputs[i].rgb, outputs[i].a));
-    const imgs = page.images.slice(0, MAX_IMG);
-    const treatments = imgs.map((im, j) => {
-      let best = 0;
-      for (let c = 1; c < 5; c++) if (out.img_logits[j * 5 + c] > out.img_logits[j * 5 + best]) best = c;
-      const t = TREAT[best];
-      const u = im.underTok;
-      return root.NadirSafeTreatment ? root.NadirSafeTreatment(t, im, u != null && u >= 0 && u < n ? outputs[u].L : null) : t;
-    });
-    return {colors, keep: [...keep], treatments};
+    return {colors, keep: [...keep], Ls: outputs.map((o) => +o.L.toFixed(3))}; // Ls: the lightness each colour is shown at
   }
 
-  const api = {buildInputs, decode, oklab, labToRgb, TREAT, MAX_TOK, MAX_IMG};
+  const api = {buildInputs, decode, buildElemInputs, decodeElements, pageElementItems, oklab, labToRgb, TREAT, IMG_ROLES, MAX_TOK, MAX_IMG, FI};
   root.NadirInfer = api;
   if (typeof module !== 'undefined') module.exports = api;
 })(typeof globalThis !== 'undefined' ? globalThis : this);

runtime/nadir-map.js

@@ -120,51 +120,15 @@
     dim: 'filter: brightness(0.88) !important;',
     plate: 'background-color: rgba(255,255,255,0.92) !important; box-shadow: 0 0 0 3px rgba(255,255,255,0.92) !important; border-radius: 4px !important;',
   };
-  // Content images must never be colour-altered: only flat graphics (logos, icons, diagrams) may get a treatment.
-  // a logo on its own white rectangle: adapt it instead of leaving a white box. Only images the page marks as its
-  // logo (logo/brand/wordmark naming, or the link home); product shots, photos and other images are never touched
-  // Returns 'adapt:<fallback>': the page tries the adaptation (its pixel check confirms the white background)
-  // and uses the fallback treatment when there is nothing to adapt. Without pixel stats yet, logos still try.
-  function adaptable(im, s, surfL) {
-    if (!(im.logo || (im.homeLink && im.header)) || im.h > 200 || (surfL != null && surfL >= 0.4)) return false;
-    if (!s) return im.w * im.h <= 250000;
-    return s.transp < 0.05 && (s.ringA || 0) >= 0.95 && (s.ringL || 0) >= 0.85 && (s.white || 0) >= 0.3 && s.colors < 0.2;
-  }
-  function baseTreatment(t, im, surfL) {
-    const s = im && im.stats;
-    // a light, flat CSS background image (texture, pale wash, pattern) sits behind content whose text is now light:
-    // clear it so the dark surface designed for that element shows through. Photos and detailed images stay.
-    // (almost no lightness variation, so never a photo; SVG backgrounds can't be photos, so a smooth gradient qualifies).
-    // Logos drawn as backgrounds are never cleared.
-    if (im && im.kind === 'bg' && !im.logo && !im.homeLink && s && s.transp < 0.05 && s.meanL > 0.8 && im.w * im.h >= 40000 &&
-        s.stdL < (/\.svg$/i.test((im.src || '').split('?')[0]) ? 0.15 : 0.07)) return 'clear';
-    // a mostly dark graphic on transparency would vanish on a dark surface: hue-preserving inversion
-    if ((!t || t === 'keep') && s && s.transp > 0.15 && s.meanL < 0.35 && im.w * im.h < 200000 && (surfL == null || surfL < 0.4)) return 'invert';
-    if (!t || t === 'keep') return 'keep';
-    const photoLike = s ? (s.transp < 0.02 && s.colors > 0.18 && s.stdL > 0.12) : (im && im.w * im.h > 60000);
-    if (photoLike && t !== 'dim') return 'keep';
-    if (t === 'dim' && im && im.w * im.h < 20000) return 'keep';
-    // a transparent graphic that is already light reads fine on a dark surface: leave it alone
-    if ((t === 'plate' || t === 'light' || t === 'invert') && s && s.transp > 0.15 && s.meanL >= 0.55 && (s.black || 0) < 0.15 &&
-        (surfL == null || surfL < 0.4)) return 'keep';
-    // no light boxes on a dark page: a dark logo on transparency gets a hue-preserving inversion instead of a plate,
-    // an opaque one already carries its own background
-    if (t === 'plate') return (!s || s.transp > 0.15) ? 'invert' : 'keep';
-    return t;
-  }
-  function safeTreatment(t, im, surfL) {
-    const base = baseTreatment(t, im, surfL);
-    return base !== 'clear' && im && adaptable(im, im.stats, surfL) ? 'adapt:' + base : base;
-  }
-
   // ---------- flat graphics on a white background (logos, icons, diagrams): adapt the image itself ----------
   // The white background becomes transparent so the page's own dark surface shows through, neutral ink flips to
   // light, and brand colours keep their hue and are lifted only as far as needed to read on dark.
   function adaptPixels(d, w, h, info) {
     info = info || {};
-    let n = 0, light = 0, sr = 0, sg = 0, sb = 0;
+    let n = 0, light = 0, sr = 0, sg = 0, sb = 0, ringN = 0;
     const ringPx = (p) => {
       const i = p * 4;
+      ringN++;
       if (d[i + 3] < 250) return;
       n++;
       const o = oklab(d[i], d[i + 1], d[i + 2]);
@@ -173,8 +137,22 @@
     for (let x = 0; x < w; x++) { ringPx(x); ringPx((h - 1) * w + x); }
     for (let y = 1; y < h - 1; y++) { ringPx(y * w); ringPx(y * w + w - 1); }
     info.ring = n ? light / n : 0;
-    if (!n || light / n < 0.5) { info.why = 'ring'; return false; } // no light background to remove (tight crops touch the edge)
-    const bg = oklab(sr / light, sg / light, sb / light);
+    let bg = null;
+    if (n && light / n >= 0.5) bg = oklab(sr / light, sg / light, sb / light);
+    else if (n > 0.2 * ringN) { info.why = 'ring'; return false; } // no light background to remove (tight crops touch the edge)
+    else {
+      // on transparency: only worth it when neutral dark ink (wordmark text) would vanish on a dark surface
+      let opq = 0, ink = 0;
+      for (let p = 0, P = w * h; p < P; p++) {
+        const i = p * 4;
+        if (d[i + 3] < 128) continue;
+        opq++;
+        const o = oklab(d[i], d[i + 1], d[i + 2]);
+        if (o[0] < 0.5 && Math.hypot(o[1], o[2]) < 0.05) ink++;
+      }
+      info.ink = opq ? ink / opq : 0;
+      if (!opq || ink / opq < 0.08) { info.why = 'no-ink'; return false; }
+    }
     const cache = new Map();
     for (let p = 0, P = w * h; p < P; p++) {
       const i = p * 4;
@@ -183,12 +161,12 @@
       let r = cache.get(key);
       if (!r) {
         const o = oklab(d[i], d[i + 1], d[i + 2]);
-        const dist = Math.hypot(o[0] - bg[0], o[1] - bg[1], o[2] - bg[2]);
-        const cov = Math.max(0, Math.min(1, (dist - 0.03) / 0.22)); // 0 = background, 1 = ink, between = anti-aliased edge
+        const dist = bg ? Math.hypot(o[0] - bg[0], o[1] - bg[1], o[2] - bg[2]) : 1;
+        const cov = bg ? Math.max(0, Math.min(1, (dist - 0.03) / 0.22)) : 1; // 0 = background, 1 = ink, between = anti-aliased edge
         if (cov <= 0) r = [0, 0, 0, 0];
         else {
           // unmix the edge pixel from the background, then restyle the ink for a dark surface
-          let L = (o[0] - (1 - cov) * bg[0]) / cov, a = (o[1] - (1 - cov) * bg[1]) / cov, b = (o[2] - (1 - cov) * bg[2]) / cov;
+          let L = bg ? (o[0] - (1 - cov) * bg[0]) / cov : o[0], a = bg ? (o[1] - (1 - cov) * bg[1]) / cov : o[1], b = bg ? (o[2] - (1 - cov) * bg[2]) / cov : o[2];
           L = Math.max(0, Math.min(1, L));
           const C = Math.hypot(a, b);
           if (C < 0.05) L = 0.94 - 0.75 * L;                       // black/gray ink -> light
@@ -240,5 +218,4 @@
   globalThis.NadirMap = NadirMap;
   globalThis.NadirAdaptImage = adaptImage;
   globalThis.NadirImageCSS = TREAT_CSS;
-  globalThis.NadirSafeTreatment = safeTreatment;
 })();

training/build_ds.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
-"""Build Nadir training data: page features + per-token dark targets (Haiku designs and native dark-mode pairs)."""
+"""Build Nadir training data: page features + per-token dark targets (Haiku designs and native dark-mode pairs),
+plus element-head targets (how each image should render, and what it is) from Haiku's contact-sheet labels."""
 import argparse, collections, difflib, gzip, hashlib, json, math, os, random, sys
 from pathlib import Path
 import numpy as np
@@ -7,8 +8,10 @@ import numpy as np
 import safety
 
 ROLES = ["bg", "text", "border", "shadow", "gradient", "fill", "stroke"]
-TREAT = ["keep", "dim", "invert", "light", "plate"]
-MAX_TOK, MAX_IMG = 256, 32
+TREAT = ["keep", "invert", "adapt", "remove"]
+LETTER = {"A": 0, "B": 1, "C": 2, "D": 3}
+IMG_ROLES = ["photo", "illustration", "diagram", "logo", "icon", "text", "backdrop", "avatar", "ui"]
+MAX_TOK, MAX_IMG, FI = 256, 32, 41
 
 
 def lin(u):
@@ -46,6 +49,62 @@ def over(top, bottom):
     return [top[i] * a + bottom[i] * (1 - a) for i in range(3)] + [1.0]
 
 
+def pad_f(f):
+    """extract.js v5 image features had 39 entries; v6 adds kind=pseudo and sprite-crop (both 0 for v5 pages)."""
+    return list(f) + [0.0] * (FI - len(f))
+
+
+def element_arrays(L, n_tok, labels):
+    """Element-head samples for the labelled images of a page (any image, not only the 32 the trunk sees)."""
+    rows = []
+    for j, labs in sorted(labels.items()):
+        if j >= len(L["images"]):
+            continue
+        im = L["images"][j]
+        if not im.get("thumb") or len(im["thumb"]) != 4096:
+            continue
+        soft = np.zeros(len(TREAT), np.float32)
+        for lab in labs:
+            best = LETTER[lab["best"]]
+            ok = [LETTER[x] for x in lab["ok"] if x != lab["best"]]
+            soft[best] += 0.75 if ok else 1.0
+            for o in ok:
+                soft[o] += 0.25 / len(ok)
+        soft /= soft.sum()
+        role = IMG_ROLES.index(labs[0]["role"]) if labs[0]["role"] in IMG_ROLES else -1
+        u = im.get("underTok", -1)
+        rows.append((pad_f(im["f"]), np.array(im["thumb"], dtype=np.uint8).reshape(32, 32, 4), im.get("under") or [1.0, 0.0, 0.0],
+                     u if 0 <= u < n_tok else -1, soft, role, j, im.get("kind", "img"), labs[0].get("bucket", ""), len(labs)))
+    K = len(rows)
+    E = dict(E_f=np.zeros((K, FI), np.float32), E_thumb=np.zeros((K, 32, 32, 4), np.uint8), E_ul=np.zeros((K, 3), np.float32),
+             E_ut=-np.ones(K, np.int64), E_soft=np.zeros((K, len(TREAT)), np.float32), E_role=-np.ones(K, np.int64), E_idx=np.zeros(K, np.int64),
+             E_kind=[], E_bucket=[], E_nlab=np.zeros(K, np.int64))
+    for k, (f, th, ul, ut, soft, role, j, kind, bucket, nl) in enumerate(rows):
+        E["E_f"][k] = f; E["E_thumb"][k] = th; E["E_ul"][k] = ul; E["E_ut"][k] = ut; E["E_soft"][k] = soft; E["E_role"][k] = role
+        E["E_idx"][k] = j; E["E_kind"].append(kind); E["E_bucket"].append(bucket); E["E_nlab"][k] = nl
+    E["E_f"] = np.nan_to_num(E["E_f"], nan=0.0, posinf=0.0, neginf=0.0)
+    return E
+
+
+def load_labels(paths):
+    """labels.jsonl rows -> {(crawl tag, page id): {image index: [label, ...]}} (a second, independent label when present)."""
+    out = collections.defaultdict(lambda: collections.defaultdict(list))
+    n = 0
+    for p in paths:
+        for line in open(p):
+            if not line.strip():
+                continue
+            try:
+                s = json.loads(line)
+            except Exception:  # a line still being written
+                continue
+            for r in s["rows"]:
+                out[(r["crawl"], r["page"])][r["img"]].append(r)
+                n += 1
+    print("element labels", n, "on", len(out), "pages", flush=True)
+    return out
+
+
 def page_arrays(L):
     toks = L["tokens"][:MAX_TOK]
     n = len(toks)
@@ -54,7 +113,7 @@ def page_arrays(L):
     rgba = np.array([t["rgba"] for t in toks], dtype=np.float32)
     rel = np.array([r for r in L["rel"] if r[1] < n and r[2] < n], dtype=np.float32).reshape(-1, 4)
     imgs = L["images"][:MAX_IMG]
-    IF = np.array([im["f"] for im in imgs], dtype=np.float32).reshape(len(imgs), 39)
+    IF = np.array([pad_f(im["f"]) for im in imgs], dtype=np.float32).reshape(len(imgs), FI)
     thumbs = np.zeros((len(imgs), 32, 32, 4), dtype=np.uint8)
     for j, im in enumerate(imgs):
         if im.get("thumb") and len(im["thumb"]) == 4096:
@@ -78,11 +137,7 @@ def haiku_targets(L, h, arrs, qc):
         i = int(k)
         if i < n:
             tgt[i] = hex2rgba(v)
-    imgt = np.full((len(arrs["IF"]),), -1, dtype=np.int64)
-    for k, v in (h.get("images") or {}).items():
-        j = int(k)
-        if j < len(imgt) and v in TREAT:
-            imgt[j] = TREAT.index(v)
+    imgt = np.full((len(arrs["IF"]),), -1, dtype=np.int64)  # page-design image labels are superseded by the element head
     # ---- QC ----
     asked = min(n, h.get("n_tok", n))
     got = np.isfinite(tgt[:asked, 0]).sum()
@@ -143,6 +198,7 @@ def repair_contrast(arrs, tgt):
             safety.hue_guard(toks[i], out[i])
     safety.keep_invisible(toks, out, arrs["rel"].tolist(), keep)
     safety.enforce_contrast(toks, out, arrs["rel"].tolist(), keep)
+    safety.rescue_surfaces(toks, out, arrs["rel"].tolist(), keep)
     changed = 0
     for i in range(n):
         if out[i]["rgb"] != before[i]:
@@ -203,78 +259,94 @@ def native_targets(L, D, arrs, qc):
     return tgt, np.full((len(arrs["IF"]),), -1, dtype=np.int64)
 
 
+LABELS = {}
+
+
+def build_one(task):
+    """One page -> training sample (or None), with its QC outcome. Runs in a worker process."""
+    tag, cdir, kind, pid, hpath, url = task
+    qc = collections.Counter()
+    crawl = Path(cdir)
+    try:
+        L = json.load(gzip.open(crawl / "pages" / pid / "light.json.gz"))
+        arrs = page_arrays(L)
+        if kind == "haiku":
+            h = json.loads(Path(hpath).read_text())
+            if h.get("refusal") or "colors" not in h:
+                return None, qc
+            r = haiku_targets(L, h, arrs, qc)
+            if r is None:
+                return None, qc
+            arrs["tgt"], arrs["imgt"] = r
+            if repair_contrast(arrs, arrs["tgt"]):
+                qc["repaired"] += 1
+            src = 0
+            url = url or L["url"]
+        else:
+            D = json.load(gzip.open(crawl / "pages" / pid / "dark.json.gz"))
+            res = native_targets(L, D, arrs, qc)
+            if res is None:
+                return None, qc
+            arrs["tgt"], arrs["imgt"] = res
+            src = 1
+        arrs["src"] = src
+        arrs["id"] = pid
+        arrs["crawl"] = tag
+        arrs["host"] = url.split("/")[2]
+        arrs.update(element_arrays(L, len(arrs["role"]), LABELS.get((tag, pid), {})))
+        return arrs, qc
+    except Exception:
+        qc["error"] += 1
+        return None, qc
+
+
 def main():
+    global LABELS
     ap = argparse.ArgumentParser()
-    ap.add_argument("--crawl", required=True)
-    ap.add_argument("--haiku", nargs="*", default=[])
+    ap.add_argument("--src", action="append", required=True, help="TAG=CRAWL_DIR=HAIKU_DIR (repeatable; HAIKU_DIR may be empty)")
+    ap.add_argument("--labels", nargs="*", default=[], help="element label files (elem_label.py labels.jsonl)")
     ap.add_argument("--out", required=True)
     ap.add_argument("--val-frac", type=float, default=0.05)
+    ap.add_argument("--workers", type=int, default=16)
     args = ap.parse_args()
-    crawl = Path(args.crawl)
-    index = {}
-    for line in (l for ip in sorted(crawl.glob("index*.jsonl")) for l in open(ip)):
-        try:
-            r = json.loads(line)
-            index[r["id"]] = r
-        except Exception:
-            pass
-    samples = []
-    qc_h = collections.Counter(); qc_n = collections.Counter()
-    # Haiku-designed pages (source 0)
-    for hd in args.haiku:
-        for fp in Path(hd).glob("*.json"):
-            if fp.name in ("ledger.jsonl",):
-                continue
-            try:
-                h = json.loads(fp.read_text())
-            except Exception:
-                continue
-            if h.get("refusal") or "colors" not in h:
-                continue
-            pid = h["id"]
+    LABELS = {k: dict(v) for k, v in load_labels(args.labels).items()}
+    tasks = []
+    for spec in args.src:
+        tag, cdir, hdir = spec.split("=")
+        crawl = Path(cdir)
+        index = {}
+        for line in (l for ip in sorted(crawl.glob("index*.jsonl")) for l in open(ip)):
             try:
-                L = json.load(gzip.open(crawl / "pages" / pid / "light.json.gz"))
+                r = json.loads(line)
+                index[r["id"]] = r
             except Exception:
-                continue
-            arrs = page_arrays(L)
-            r = haiku_targets(L, h, arrs, qc_h)
-            if r is None:
-                continue
-            arrs["tgt"], arrs["imgt"] = r
-            if repair_contrast(arrs, arrs["tgt"]):
-                qc_h["repaired"] += 1
-            arrs["src"] = 0
-            arrs["id"] = pid
-            arrs["host"] = index.get(pid, {}).get("url", L["url"]).split("/")[2]
-            samples.append(arrs)
-    # native dark-mode pairs (source 1)
-    for pid, r in index.items():
-        if r.get("kind") not in ("native_media", "native_class"):
-            continue
-        try:
-            L = json.load(gzip.open(crawl / "pages" / pid / "light.json.gz"))
-            D = json.load(gzip.open(crawl / "pages" / pid / "dark.json.gz"))
-        except Exception:
-            continue
-        arrs = page_arrays(L)
-        res = native_targets(L, D, arrs, qc_n)
-        if res is None:
-            continue
-        arrs["tgt"], arrs["imgt"] = res
-        arrs["src"] = 1
-        arrs["id"] = pid
-        arrs["host"] = r["url"].split("/")[2]
-        samples.append(arrs)
-    print("haiku QC", dict(qc_h), "native QC", dict(qc_n), "samples", len(samples), flush=True)
+                pass
+        for fp in (Path(hdir).glob("*.json") if hdir else []):
+            if len(fp.stem) == 16:
+                tasks.append((tag, cdir, "haiku", fp.stem, str(fp), index.get(fp.stem, {}).get("url")))
+        for pid, r in index.items():
+            if r.get("kind") in ("native_media", "native_class"):
+                tasks.append((tag, cdir, "native", pid, None, r["url"]))
+    print(len(tasks), "pages to build", flush=True)
+    import multiprocessing as mp
+    samples, qc = [], collections.Counter()
+    with mp.get_context("fork").Pool(args.workers) as pool:
+        for arrs, q in pool.imap_unordered(build_one, tasks, chunksize=32):
+            qc.update(q)
+            if arrs is not None:
+                samples.append(arrs)
+    samples.sort(key=lambda s: (s["crawl"], s["src"], s["id"]))
+    print("QC", dict(qc), "samples", len(samples), "elements", sum(len(s["E_ut"]) for s in samples), flush=True)
     hosts = sorted(set(s["host"].replace("www.", "") for s in samples))
-    # stable split: a host's side never changes as the crawl grows
+    # stable split: a host's side never changes as the crawl grows (and a site's inner pages share its side)
     val_hosts = {h for h in hosts if int(hashlib.md5(h.encode()).hexdigest()[:8], 16) % 1000 < args.val_frac * 1000}
     train = [s for s in samples if s["host"].replace("www.", "") not in val_hosts]
     val = [s for s in samples if s["host"].replace("www.", "") in val_hosts]
     import pickle
     with open(args.out, "wb") as f:
-        pickle.dump({"train": train, "val": val, "roles": ROLES, "treat": TREAT}, f, protocol=4)
-    print("train", len(train), "val", len(val), "haiku", sum(s["src"] == 0 for s in samples), "native", sum(s["src"] == 1 for s in samples))
+        pickle.dump({"train": train, "val": val, "roles": ROLES, "treat": TREAT, "img_roles": IMG_ROLES}, f, protocol=4)
+    print("train", len(train), "val", len(val), "haiku", sum(s["src"] == 0 for s in samples), "native", sum(s["src"] == 1 for s in samples),
+          "elements train", sum(len(s["E_ut"]) for s in train), "val", sum(len(s["E_ut"]) for s in val))
 
 
 if __name__ == "__main__":

training/crawl.py

@@ -0,0 +1,276 @@
+#!/usr/bin/env python3
+"""Nadir crawler: render sites light + emulated dark, extract features, classify dark-mode support.
+
+Output per page: OUT/pages/<id>/{light.json.gz, dark.json.gz?, shot.jpg, meta.json}
+and a line in OUT/index.jsonl.
+"""
+import argparse, asyncio, base64, gzip, hashlib, json, os, random, re, sys, time, traceback
+from pathlib import Path
+from playwright.async_api import async_playwright
+
+HERE = Path(__file__).resolve().parent
+EXTRACT_JS = (HERE / "extract.js").read_text()
+
+NO_ANIM = """(() => { try { const s = new CSSStyleSheet(); s.replaceSync('*,*::before,*::after{transition:none!important;animation-duration:0s!important;animation-delay:0s!important;caret-color:auto}');
+ document.adoptedStyleSheets = [...document.adoptedStyleSheets, s]; window.__nadirNoAnim = s; } catch (e) {} })()"""
+DROP_NO_ANIM = """(() => { try { document.adoptedStyleSheets = document.adoptedStyleSheets.filter(x => x !== window.__nadirNoAnim); } catch (e) {} })()"""
+
+FETCH_SHIM = """window.__nadirFetchImage = async (src) => {
+  try { const r = await window.__nadirFetchB64(src); if (!r) return null;
+    const bin = atob(r.b64); const u = new Uint8Array(bin.length); for (let i = 0; i < bin.length; i++) u[i] = bin.charCodeAt(i);
+    return new Blob([u], {type: r.ct || 'application/octet-stream'}); } catch (e) { return null; } };"""
+
+
+def page_id(url):
+    return hashlib.sha1(url.encode()).hexdigest()[:16]
+
+
+async def settle(page, ms=600):
+    try:
+        await page.wait_for_timeout(ms)
+    except Exception:
+        pass
+
+
+async def autoscroll(page):
+    # trigger lazy content, then return to top
+    try:
+        await page.evaluate("""async () => {
+          const H = Math.min(document.documentElement.scrollHeight, innerHeight * 6);
+          for (let y = 0; y < H; y += innerHeight * 0.8) { scrollTo(0, y); await new Promise(r => setTimeout(r, 120)); }
+          scrollTo(0, 0); await new Promise(r => setTimeout(r, 250)); }""")
+    except Exception:
+        pass
+
+
+async def extract(page, **opts):
+    await page.evaluate(EXTRACT_JS)
+    o = dict(opts)
+    js = """async (o) => await __nadirExtract(o)"""
+    return await page.evaluate(js, o)
+
+
+DARK_TOGGLE_JS = """(sel) => {
+  const html = document.documentElement, body = document.body;
+  const undo = [];
+  const m = sel.match(/^(html|body|:root)?(.*)$/);
+  const target = (m[1] === 'body') ? body : html;
+  const rest = m[2];
+  let cm = rest.match(/^\\.([\\w-]+)$/);
+  let am = rest.match(/^\\[([\\w-]+)(?:[~|^*$]?=\\s*["']?([\\w-]+)["']?)?\\s*(?:i\\s*)?\\]$/);
+  if (cm) { for (const t of [html, body]) { if (!t.classList.contains(cm[1])) { t.classList.add(cm[1]); undo.push(['c', t, cm[1]]); } } }
+  else if (am) { const v = am[2] || 'dark'; for (const t of [html, body]) { const old = t.getAttribute(am[1]); t.setAttribute(am[1], v); undo.push(['a', t, am[1], old]); } }
+  else return false;
+  window.__nadirUndo = undo; return true; }"""
+UNDO_JS = """() => { for (const u of (window.__nadirUndo || [])) { if (u[0] === 'c') u[1].classList.remove(u[2]); else if (u[3] === null) u[1].removeAttribute(u[2]); else u[1].setAttribute(u[2], u[3]); } }"""
+
+
+LINKS_JS = r"""() => { const o = location.origin, seen = new Set(), out = [];
+  for (const a of document.querySelectorAll('a[href]')) { let u; try { u = new URL(a.href); } catch (e) { continue; }
+    if (u.origin !== o || u.pathname.length < 2 || /\.(pdf|zip|png|jpe?g|gif|svg|mp4|xml|json)$/i.test(u.pathname) || /login|signin|signup|register|cart|checkout|account|logout|auth/i.test(u.pathname)) continue;
+    u.hash = ''; const k = u.pathname; if (seen.has(k)) continue; seen.add(k); out.push(u.href); if (out.length >= 12) break; }
+  return out; }"""
+
+
+IMG_STATS_JS = """async (a) => {
+  const res = await __nadirImageStats(a.images, {shot: a.shot, bigThumbSize: 96, concurrency: 8});
+  a.images.forEach((o, i) => { Object.assign(o, res[i]); o.f = __nadirImageFeatures(o, a.vw, a.vh, a.docH); });
+  return a.images; }"""
+
+
+async def image_stats(ctx, d, shot_png):
+    """Pixel stats in a separate blank page (like the extension's offscreen document): no page CSP, no Python transfer."""
+    if not d["images"]:
+        return
+    helper = await ctx.new_page()
+    try:
+        await helper.evaluate(EXTRACT_JS)
+        shot = "data:image/png;base64," + base64.b64encode(shot_png).decode()
+        d["images"] = await asyncio.wait_for(helper.evaluate(IMG_STATS_JS, {"images": d["images"], "shot": shot, "vw": d["vw"], "vh": d["vh"], "docH": d["docH"]}), 40)
+    except Exception as e:
+        d["imgStatsError"] = str(e)[:200]
+    finally:
+        await helper.close()
+
+
+def is_dark(d):
+    dk = d["darkness"]
+    return dk["bgL"] < 0.45 and dk["darkFrac"] > 0.6 and (dk["textL"] == 0 or dk["textL"] > 0.55)
+
+
+async def crawl_one(ctx, url, out_dir, args):
+    pid = page_id(url)
+    pdir = out_dir / "pages" / pid
+    rec = {"id": pid, "url": url, "t": time.time()}
+    page = await ctx.new_page()
+    try:
+        t0 = time.time()
+        resp = await page.goto(url, wait_until="domcontentloaded", timeout=args.nav_timeout * 1000)
+        rec["status"] = resp.status if resp else None
+        rec["final_url"] = page.url
+        try:
+            await page.wait_for_load_state("load", timeout=8000)
+        except Exception:
+            pass
+        await settle(page, 1200)
+        await page.evaluate(NO_ANIM)
+        await autoscroll(page)
+        await settle(page, 700)
+        light = await extract(page, describe=True, thumbs=False, elemColors=True, markSvg=True, maxImages=args.max_images)
+        shot_png = await page.screenshot(type="png")
+        await image_stats(ctx, light, shot_png)
+        rec["light_ms"] = light["stats"]["ms"]
+        rec["n_tokens"] = len(light["tokens"])
+        rec["n_images"] = len(light["images"])
+        rec["darkness"] = light["darkness"]
+        rec["scheme"] = light["scheme"]
+        if light["stats"]["visible"] < 15 or len(light["tokens"]) < 4:
+            rec["kind"] = "empty"
+            return rec
+        pdir.mkdir(parents=True, exist_ok=True)
+        # screenshot (top ~2 screens)
+        H = min(int(light["docH"]), args.shot_h)
+        await page.screenshot(path=str(pdir / "shot.jpg"), type="jpeg", quality=80, full_page=True,
+                              clip={"x": 0, "y": 0, "width": 1280, "height": H})
+        light_dark = is_dark(light)
+        kind = "dark" if light_dark else "light"
+        dark = None
+        # 1) emulated prefers-color-scheme: dark
+        await page.emulate_media(color_scheme="dark")
+        await settle(page, 1500)
+        dm = await extract(page, describe=False, thumbs=False, elemColors=True, pseudo=True)
+        if not light_dark and is_dark(dm):
+            kind, dark = "native_media", dm
+        elif light_dark:
+            pass
+        else:
+            # 2) class / attribute toggles found in stylesheets
+            await page.emulate_media(color_scheme="light")
+            await settle(page, 400)
+            for sel, n in (light["scheme"]["darkSelectors"] or [])[:3]:
+                if n < 8:
+                    continue
+                ok = await page.evaluate(DARK_TOGGLE_JS, sel)
+                if not ok:
+                    continue
+                await settle(page, 900)
+                dc = await extract(page, describe=False, thumbs=False, elemColors=True)
+                if is_dark(dc):
+                    kind, dark = "native_class", dc
+                    rec["toggle"] = sel
+                    break
+                await page.evaluate(UNDO_JS)
+                await settle(page, 300)
+        rec["kind"] = kind
+        try:
+            rec["links"] = await page.evaluate(LINKS_JS)
+        except Exception:
+            rec["links"] = []
+        if dark is not None:
+            rec["dark_darkness"] = dark["darkness"]
+            await page.screenshot(path=str(pdir / "shot_dark.jpg"), type="jpeg", quality=80, full_page=True,
+                                  clip={"x": 0, "y": 0, "width": 1280, "height": H})
+            with gzip.open(pdir / "dark.json.gz", "wt") as f:
+                json.dump(dark, f, separators=(",", ":"))
+        with gzip.open(pdir / "light.json.gz", "wt") as f:
+            json.dump(light, f, separators=(",", ":"))
+        rec["secs"] = round(time.time() - t0, 1)
+        return rec
+    except Exception as e:
+        rec["kind"] = "error"
+        rec["error"] = (type(e).__name__ + ": " + str(e))[:300]
+        return rec
+    finally:
+        try:
+            await page.close()
+        except Exception:
+            pass
+
+
+async def worker(wid, browser_pool, queue, out_dir, args, index_f, counters):
+    while True:
+        try:
+            url = queue.get_nowait()
+        except asyncio.QueueEmpty:
+            return
+        browser = browser_pool[wid % len(browser_pool)]
+        ctx = None
+        try:
+            ctx = await browser.new_context(viewport={"width": 1280, "height": 800}, device_scale_factor=1,
+                                            color_scheme="light", locale="en-US", ignore_https_errors=True,
+                                            user_agent=args.ua, java_script_enabled=True, bypass_csp=True)
+            ctx.set_default_timeout(15000)
+            rec = await asyncio.wait_for(crawl_one(ctx, url, out_dir, args), timeout=args.page_timeout)
+        except asyncio.TimeoutError:
+            rec = {"id": page_id(url), "url": url, "kind": "error", "error": "page timeout"}
+        except Exception as e:
+            rec = {"id": page_id(url), "url": url, "kind": "error", "error": str(e)[:300]}
+        finally:
+            if ctx:
+                try:
+                    await asyncio.wait_for(ctx.close(), 10)
+                except Exception:
+                    pass
+        counters[rec.get("kind", "error")] = counters.get(rec.get("kind", "error"), 0) + 1
+        index_f.write(json.dumps(rec) + "\n")
+        index_f.flush()
+
+
+async def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--urls", required=True)
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--workers", type=int, default=24)
+    ap.add_argument("--browsers", type=int, default=4)
+    ap.add_argument("--nav-timeout", type=int, default=25)
+    ap.add_argument("--page-timeout", type=int, default=90)
+    ap.add_argument("--shot-h", type=int, default=2000)
+    ap.add_argument("--max-images", type=int, default=48, help="images kept per page (the model sees the first 48; round 2 keeps more for the element head)")
+    ap.add_argument("--limit", type=int, default=0)
+    ap.add_argument("--shard", type=int, default=0)
+    ap.add_argument("--nshards", type=int, default=1)
+    ap.add_argument("--ua", default="Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/141.0.0.0 Safari/537.36")
+    args = ap.parse_args()
+    out_dir = Path(args.out)
+    (out_dir / "pages").mkdir(parents=True, exist_ok=True)
+    done = set()
+    for ip in out_dir.glob("index*.jsonl"):
+        for line in ip.open():
+            try:
+                done.add(json.loads(line)["url"])
+            except Exception:
+                pass
+    idx = out_dir / ("index.jsonl" if args.nshards == 1 else f"index.{args.shard}.jsonl")
+    urls = [u.strip() for u in open(args.urls) if u.strip() and not u.startswith("#")]
+    urls = [u for u in urls if u not in done and (args.nshards == 1 or int(hashlib.sha1(u.encode()).hexdigest(), 16) % args.nshards == args.shard)]
+    if args.limit:
+        urls = urls[: args.limit]
+    print(f"{len(urls)} urls to crawl ({len(done)} done)", flush=True)
+    q = asyncio.Queue()
+    for u in urls:
+        q.put_nowait(u)
+    counters = {}
+    async with async_playwright() as p:
+        browsers = []
+        for _ in range(args.browsers):
+            browsers.append(await p.chromium.launch(headless=True, args=[
+                "--disable-web-security", "--disable-features=IsolateOrigins,site-per-process,BlockInsecurePrivateNetworkRequests",
+                "--disable-dev-shm-usage", "--no-first-run", "--mute-audio", "--hide-scrollbars", "--font-render-hinting=none",
+                "--autoplay-policy=user-gesture-required", "--disable-notifications"]))
+        with idx.open("a") as f:
+            start = time.time()
+            async def report():
+                while True:
+                    await asyncio.sleep(60)
+                    n = sum(counters.values())
+                    print(f"[{time.strftime('%H:%M:%S')}] {n} done, {n / max(1, time.time() - start) * 3600:.0f}/h {counters}", flush=True)
+            rep = asyncio.create_task(report())
+            await asyncio.gather(*[worker(i, browsers, q, out_dir, args, f, counters) for i in range(args.workers)])
+            rep.cancel()
+        for b in browsers:
+            await b.close()
+    print("DONE", counters, flush=True)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

training/elem_eval.py

@@ -0,0 +1,110 @@
+#!/usr/bin/env python3
+"""Element-head evaluation on held-out sites: does each image get a rendering Haiku judged acceptable?
+
+Compares the v2 element head with what v1 shipped for the same images (its 5-class image head for the first 32
+images of a page, the hand-written runtime rules for everything else), per kind of image.
+"""
+import argparse, collections, gzip, importlib.util, json, pickle, sys
+from pathlib import Path
+import numpy as np
+import torch
+
+from model import NadirNet, TREAT
+from train import collate, prep, element_logits
+import safety
+
+V1_MAP = {"keep": 0, "dim": 0, "invert": 1, "light": 1, "plate": 1, "clear": 3}
+
+
+def v1_decision(t, im, surf_L):
+    """v1 runtime: model treatment (or 'keep' for images it never saw) through the hand-written rules."""
+    s = safety.safe_treatment(t, im, surf_L)
+    if s.startswith("adapt:"):
+        return 2  # the logo adaptation path
+    return V1_MAP.get(s, 0)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ckpt", required=True)
+    ap.add_argument("--data", required=True)
+    ap.add_argument("--crawl", action="append", required=True, help="TAG=CRAWL_DIR")
+    ap.add_argument("--v1-ckpt", default="")
+    ap.add_argument("--v1-src", default="", help="v1 model.py (its architecture)")
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--rows", default="", help="also write every element (probabilities, labels) as jsonl")
+    args = ap.parse_args()
+    crawls = dict(c.split("=", 1) for c in args.crawl)
+    dev = "cuda" if torch.cuda.is_available() else "cpu"
+    ck = torch.load(args.ckpt, map_location="cpu")
+    T = float(ck.get("temperature", 1.0))
+    m = NadirNet(**ck["cfg"]); m.load_state_dict(ck["model"]); m.eval().to(dev)
+    v1 = None
+    if args.v1_ckpt and args.v1_src:
+        spec = importlib.util.spec_from_file_location("model_v1", args.v1_src)
+        mv1 = importlib.util.module_from_spec(spec); spec.loader.exec_module(mv1)
+        c1 = torch.load(args.v1_ckpt, map_location="cpu")
+        v1 = mv1.NadirNet(**c1["cfg"]); v1.load_state_dict(c1["model"]); v1.eval().to(dev)
+    data = pickle.load(open(args.data, "rb"))
+    val = [prep(s) for s in data["val"] if len(s.get("E_ut", []))]
+    rows = []
+    with torch.no_grad():
+        for s in val:
+            b = collate([s], dev)
+            lab, _, ctx = m(b["tok_f"], b["role"], b["lab"], b["img_f"], b["thumbs"], b["page"], b["src"], b["rel"], b["mask"])
+            tl, rl = element_logits(m, b, lab.float(), ctx)
+            k = int(b["e_mask"][0].sum())
+            p = torch.softmax(tl[0, :k] / T, -1).cpu().numpy()
+            role = rl[0, :k].argmax(-1).cpu().numpy()
+            rprob = torch.softmax(rl[0, :k].float() / T, -1).max(-1).values.cpu().numpy()
+            v1t = None
+            if v1 is not None and len(s["IF"]):
+                _, _, il = v1(b["tok_f"], b["role"], b["lab"], b["img_f"][..., :39], b["thumbs"], b["page"], b["src"], b["rel"], b["mask"])
+                v1t = il[0].float().argmax(-1).cpu().numpy()
+            L = json.load(gzip.open(Path(crawls[s["crawl"]]) / "pages" / s["id"] / "light.json.gz"))
+            n = len(s["role"])
+            outL = lab[0, :n, 0].float().cpu().numpy()
+            for q in range(k):
+                j = int(s["E_idx"][q])
+                im = L["images"][j]
+                u = im.get("underTok", -1)
+                surf_L = float(outL[u]) if 0 <= u < n else None
+                ud = [round(float(x), 4) for x in lab[0, u].float().cpu().numpy()] if 0 <= u < n else None
+                old = ["keep", "dim", "invert", "light", "plate"][int(v1t[j])] if (v1t is not None and j < 32 and j < len(v1t)) else "keep"
+                soft = s["E_soft"][q]
+                rows.append({"bucket": s["E_bucket"][q], "kind": s["E_kind"][q], "soft": soft.tolist(), "best": int(np.argmax(soft)),
+                             "v2": int(np.argmax(p[q])), "conf": float(p[q].max()), "p": [round(float(x), 4) for x in p[q]], "role_pred": int(role[q]), "role": int(s["E_role"][q]),
+                             "v1": v1_decision(old, im, surf_L), "v1_seen": bool(v1t is not None and j < 32), "nlab": int(s["E_nlab"][q]),
+                             **({"rp": round(float(rprob[q]), 4), "f": [round(float(x), 4) for x in s["E_f"][q]], "ud": ud, "ul": [round(float(x), 4) for x in s["E_ul"][q]]} if args.rows else {})})
+    def score(rs, key):
+        if not rs:
+            return {}
+        acc = np.mean([r[key] == r["best"] for r in rs])
+        ok = np.mean([r["soft"][r[key]] > 0 for r in rs])
+        return {"n": len(rs), "acc": round(float(acc), 4), "acceptable": round(float(ok), 4)}
+    rep = {"temperature": T, "overall": {"v2": score(rows, "v2"), "v1": score(rows, "v1")}, "by_bucket": {}, "by_kind": {}}
+    for key in ("bucket", "kind"):
+        groups = collections.defaultdict(list)
+        for r in rows:
+            groups[r[key]].append(r)
+        rep["by_" + key] = {g: {"v2": score(rs, "v2"), "v1": score(rs, "v1")} for g, rs in sorted(groups.items())}
+    # calibration: confidence vs accuracy
+    bins = collections.defaultdict(list)
+    for r in rows:
+        bins[min(9, int(r["conf"] * 10))].append(r["v2"] == r["best"])
+    rep["calibration"] = {f"{b / 10:.1f}": {"n": len(v), "acc": round(float(np.mean(v)), 3)} for b, v in sorted(bins.items())}
+    rr = [r for r in rows if r["role"] >= 0]
+    rep["role_acc"] = round(float(np.mean([r["role_pred"] == r["role"] for r in rr])), 4) if rr else None
+    # Haiku's own consistency where a second, independent label exists
+    rep["n_double_labelled"] = sum(r["nlab"] > 1 for r in rows)
+    json.dump(rep, open(args.out, "w"), indent=1)
+    if args.rows:
+        with open(args.rows, "w") as f:
+            for r in rows:
+                f.write(json.dumps(r) + "\n")
+    print(json.dumps({k: rep[k] for k in ("overall", "role_acc", "temperature")}, indent=1))
+    print(json.dumps(rep["by_bucket"], indent=1))
+
+
+if __name__ == "__main__":
+    main()

training/elem_label.py

@@ -0,0 +1,259 @@
+#!/usr/bin/env python3
+"""Element-head teacher data, step 2: Claude Haiku 5.5 picks how each image should render in dark mode.
+
+Reads the contact sheets from elem_sheets.py and asks, per row, for the best rendering (A keep, B invert, C adapt,
+D remove), every acceptable one, and the image's role. Batch API (half price) with a cost ledger and a hard budget;
+reads ANTHROPIC_API_KEY from the environment. Output: OUT/labels.jsonl, one line per sheet.
+"""
+import argparse, base64, json, os, random, sys, time, urllib.request
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+MODEL = "claude-haiku-5-5"
+PRICE_IN, PRICE_OUT = 0.10 / 1e6, 0.50 / 1e6
+PROXY = os.environ.get("ANTHROPIC_BASE_URL", "https://api.anthropic.com")
+HDRS = {"x-api-key": os.environ.get("ANTHROPIC_API_KEY", ""), "anthropic-version": "2023-06-01", "content-type": "application/json"}
+ROLES = ["photo", "illustration", "diagram", "logo", "icon", "text", "backdrop", "avatar", "ui"]
+OPTS = ["A", "B", "C", "D"]
+
+SYSTEM = """You are a senior product designer reviewing how a dark-mode browser extension renders the images of light websites. Each numbered row of the sheet is one image from a real web page.
+
+Left: the image in its original light page (inside the red box), or on its light background when the page area was not captured. Then the candidate renderings on the dark surface that the dark theme puts behind that image. These are the exact pixels the browser would show:
+- A keep: unchanged.
+- B invert: hue-preserving inversion. Dark becomes light and light becomes dark while hues stay, so yellows, pastels and white areas turn dark.
+- C adapt: only dark neutral ink (black or gray text, lines, glyphs) becomes light, and a white background box is removed; colours are kept and lifted just enough to read on dark.
+- D remove: the background image is hidden so the dark surface shows through (offered only for large CSS backgrounds; the sample text shows how content on top of it will look in dark mode).
+"n/a" means that rendering is not possible for this image.
+
+Pick the rendering a careful designer would ship in a polished dark theme:
+- Photos, video thumbnails, artwork, product shots, screenshots, maps, avatars, flags, emoji, ads and colourful illustrations must look exactly as their author intended: A. Never alter the colours of real content.
+- Everything must stay clearly visible. Dark icons, arrows, glyphs, text rendered as an image, math formulas and line art that would vanish or be hard to see on the dark surface must become light: B or C, whichever looks more natural.
+- Logos must stay recognizable and true to the brand. Keep brand colours: A when the logo already reads well on the dark surface, C when it mixes brand colours with dark text or sits in a white box, B only if the inverted logo still looks like the brand.
+- Diagrams, charts and documents on white: A unless the inverted or adapted version stays faithful and clearly readable.
+- Large light textures, pale washes, subtle patterns and plain light panels behind content: D. Photographic or illustrated backgrounds: A.
+- A white box around an image is acceptable for content (photos, screenshots), not for icons and logos when a cleaner option exists.
+
+For every row return its number k, the image's role, the best rendering, and every rendering you would consider acceptable (always including the best). Roles: photo, illustration, diagram (charts, maps, documents, screenshots of text), logo, icon (UI glyphs and symbols), text (text or math rendered as an image), backdrop (decorative background behind content), avatar (people, profile pictures), ui (other interface decoration: dividers, bullets, shapes, badges)."""
+
+SCHEMA = {
+    "type": "object",
+    "properties": {"rows": {"type": "array", "items": {"type": "object", "properties": {
+        "k": {"type": "integer"}, "role": {"type": "string", "enum": ROLES}, "best": {"type": "string", "enum": OPTS},
+        "ok": {"type": "array", "items": {"type": "string", "enum": OPTS}}}, "required": ["k", "role", "best", "ok"], "additionalProperties": False}}},
+    "required": ["rows"], "additionalProperties": False,
+}
+
+
+def http(method, path, body=None, timeout=300):
+    req = urllib.request.Request(PROXY + path, data=json.dumps(body).encode() if body is not None else None, headers=HDRS, method=method)
+    with urllib.request.urlopen(req, timeout=timeout) as r:
+        return r.read()
+
+
+def as_jpeg(png):
+    from PIL import Image
+    import io
+    buf = io.BytesIO()
+    Image.open(io.BytesIO(png)).convert("RGB").save(buf, "JPEG", quality=88)
+    return buf.getvalue()
+
+
+def params(sheet_png, n_rows, args):
+    content = [{"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": base64.b64encode(as_jpeg(sheet_png)).decode()}},
+               {"type": "text", "text": f"This sheet has {n_rows} rows (#1-#{n_rows}). Return one entry per row."}]
+    return {"model": MODEL, "max_tokens": args.max_out, "system": SYSTEM, "thinking": {"type": "adaptive"},
+            "output_config": {"effort": args.effort, "format": {"type": "json_schema", "schema": SCHEMA}},
+            "messages": [{"role": "user", "content": content}]}
+
+
+def parse(msg, rows):
+    txt = "".join(b.get("text", "") for b in msg.get("content", []) if b.get("type") == "text")
+    ans = json.loads(txt)["rows"]
+    out = {}
+    for a in ans:
+        k = int(a["k"])
+        if not 1 <= k <= len(rows) or k in out:
+            continue
+        offer = rows[k - 1]["offer"]
+        best = a["best"] if a["best"] in offer else None
+        ok = sorted(set(x for x in a["ok"] if x in offer) | ({best} if best else set()))
+        if best:
+            out[k] = {"role": a["role"], "best": best, "ok": ok}
+    return out
+
+
+def usage_cost(u, batch=True):
+    c = u.get("input_tokens", 0) * PRICE_IN + u.get("output_tokens", 0) * PRICE_OUT
+    return c * (0.5 if batch else 1.0)
+
+
+class Ledger:
+    def __init__(self, path):
+        self.path, self.cost, self.n = path, 0.0, 0
+        if path.exists():
+            for l in path.open():
+                r = json.loads(l)
+                self.cost += r["cost"]
+                self.n += 1
+
+    def add(self, rec):
+        self.cost += rec["cost"]
+        self.n += 1
+        with self.path.open("a") as f:
+            f.write(json.dumps(rec) + "\n")
+
+
+def load_sheets(sheet_dir):
+    sheets = {}
+    for mp in sorted(Path(sheet_dir).glob("meta.*.jsonl")):
+        for l in mp.open():
+            if l.strip():
+                s = json.loads(l)
+                sheets[s["sheet"]] = s
+    return sheets
+
+
+def write_label(out_f, sid, meta, labels, rec, dup=False):
+    rows = []
+    for r in meta["rows"]:
+        lab = labels.get(r["k"])
+        if lab:
+            rows.append(dict(r, **lab))
+    out_f.write(json.dumps({"sheet": sid, "dup": dup, "type": meta["type"], "rows": rows, "usage": rec}) + "\n")
+    out_f.flush()
+
+
+def pilot(args, sheets, ledger, out_dir):
+    ids = list(sheets)[: args.pilot]
+    with (out_dir / "pilot.jsonl").open("a") as f:
+        def one(sid):
+            meta = sheets[sid]
+            png = (Path(args.sheets) / "sheets" / f"{sid}.png").read_bytes()
+            t0 = time.time()
+            m = json.loads(http("POST", "/v1/messages", params(png, len(meta["rows"]), args), timeout=600))
+            u = m.get("usage", {})
+            rec = {"id": sid, "cost": usage_cost(u, batch=False), "in": u.get("input_tokens"), "out": u.get("output_tokens"), "secs": round(time.time() - t0, 1)}
+            ledger.add(rec)
+            write_label(f, sid, meta, parse(m, meta["rows"]), rec)
+            return rec
+        with ThreadPoolExecutor(6) as ex:
+            for rec in ex.map(one, ids):
+                print(rec, flush=True)
+    print(f"pilot spent ${ledger.cost:.4f} on {ledger.n}", flush=True)
+
+
+def batch(args, sheets, ledger, out_dir):
+    state_p = out_dir / "batches.jsonl"
+    batches = {}
+    if state_p.exists():
+        for l in state_p.open():
+            b = json.loads(l)
+            batches.setdefault(b["id"], {}).update(b)
+    submitted = set(s for b in batches.values() for s in b.get("ids", []))
+    labels_f = (out_dir / "labels.jsonl").open("a")
+    est = args.est
+    rnd = random.Random(11)
+    while True:
+        pending = 0.0
+        for bid, b in list(batches.items()):
+            if b.get("collected"):
+                continue
+            try:
+                info = json.loads(http("GET", f"/v1/messages/batches/{bid}"))
+            except Exception as e:
+                print("poll error", str(e)[:200], flush=True)
+                pending += len(b["ids"]) * est
+                continue
+            if info["processing_status"] != "ended":
+                pending += len(b["ids"]) * est
+                continue
+            raw = http("GET", f"/v1/messages/batches/{bid}/results", timeout=900).decode()
+            ok = bad = 0
+            for line in raw.splitlines():
+                if not line.strip():
+                    continue
+                r = json.loads(line)
+                cid = r["custom_id"]
+                sid, dup = cid.split("_")[0], cid.endswith("_dup")
+                res = r["result"]
+                if res["type"] != "succeeded":
+                    bad += 1
+                    continue
+                m = res["message"]
+                u = m.get("usage", {})
+                rec = {"id": cid, "cost": usage_cost(u), "in": u.get("input_tokens"), "out": u.get("output_tokens"), "batch": bid}
+                ledger.add(rec)
+                try:
+                    write_label(labels_f, sid, sheets[sid], parse(m, sheets[sid]["rows"]), rec, dup)
+                    ok += 1
+                except Exception:
+                    bad += 1
+            b["collected"] = True
+            with state_p.open("a") as f:
+                f.write(json.dumps({"id": bid, "collected": True}) + "\n")
+            print(f"[{time.strftime('%H:%M:%S')}] collected {bid}: {ok} ok, {bad} failed", flush=True)
+        if ledger.n > 100:
+            est = max(0.0005, ledger.cost / ledger.n)
+        if args.follow:
+            sheets.update(load_sheets(args.sheets))
+        new = [s for s in sheets if s not in submitted]
+        room = args.budget - ledger.cost - pending
+        can = int(max(0, room) / est)
+        if new and can > 0:
+            chunk = new[: min(args.batch_size, can)]
+            reqs = []
+            for sid in chunk:
+                png = (Path(args.sheets) / "sheets" / f"{sid}.png").read_bytes()
+                p = params(png, len(sheets[sid]["rows"]), args)
+                reqs.append({"custom_id": sid, "params": p})
+                if rnd.random() < args.dup:  # independent second opinion: measures label agreement
+                    reqs.append({"custom_id": sid + "_dup", "params": p})
+            try:
+                info = json.loads(http("POST", "/v1/messages/batches", {"requests": reqs}, timeout=900))
+                b = {"id": info["id"], "ids": chunk, "t": time.time()}
+                batches[info["id"]] = b
+                with state_p.open("a") as f:
+                    f.write(json.dumps(b) + "\n")
+                submitted.update(chunk)
+                print(f"[{time.strftime('%H:%M:%S')}] submitted {info['id']} with {len(reqs)} (spent ${ledger.cost:.2f}, pending ~${pending:.2f}, est/sheet ${est:.4f})", flush=True)
+                continue
+            except Exception as e:
+                print("submit error", str(e)[:300], flush=True)
+                time.sleep(60)
+        open_b = [b for b in batches.values() if not b.get("collected")]
+        print(f"[{time.strftime('%H:%M:%S')}] spent ${ledger.cost:.3f} on {ledger.n} | open batches {len(open_b)} | unsubmitted {len(new)}", flush=True)
+        if not open_b and (not new or can <= 0) and not args.follow:
+            break
+        if room <= 0 and not open_b:
+            print("budget reached", flush=True)
+            break
+        time.sleep(args.poll)
+    print(f"DONE ${ledger.cost:.3f} {ledger.n}", flush=True)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--sheets", required=True)
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--budget", type=float, default=25.0)
+    ap.add_argument("--effort", default="high")
+    ap.add_argument("--max-out", type=int, default=16000)
+    ap.add_argument("--pilot", type=int, default=0, help="label N sheets with the realtime API and stop")
+    ap.add_argument("--batch-size", type=int, default=300)
+    ap.add_argument("--dup", type=float, default=0.02)
+    ap.add_argument("--est", type=float, default=0.002)
+    ap.add_argument("--poll", type=int, default=60)
+    ap.add_argument("--follow", action="store_true", help="keep picking up new sheets")
+    args = ap.parse_args()
+    out_dir = Path(args.out)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    sheets = load_sheets(args.sheets)
+    ledger = Ledger(out_dir / ("ledger_pilot.jsonl" if args.pilot else "ledger.jsonl"))
+    print(len(sheets), "sheets", flush=True)
+    if args.pilot:
+        return pilot(args, sheets, ledger, out_dir)
+    batch(args, sheets, ledger, out_dir)
+
+
+if __name__ == "__main__":
+    main()

training/elem_sheets.py

@@ -0,0 +1,408 @@
+#!/usr/bin/env python3
+"""Element-head teacher data, step 1: contact sheets showing each page image every way the extension can render it.
+
+For each sampled image (logo, icon, formula, photo, backdrop...) of a crawled light page that has a Haiku dark design,
+one row: the element in its light page (red outline), then the candidates on the dark surface designed for that spot:
+  A keep    - unchanged
+  B invert  - exactly Chrome's `invert(1) hue-rotate(180deg)` (dark <-> light, hues kept)
+  C adapt   - the runtime's adaptPixels: neutral dark ink -> light, white box removed, brand colours kept
+  D remove  - large CSS background images only: the designed dark surface shows through
+Writes OUT/sheets/<id>.png and OUT/meta.<worker>.jsonl (one line per sheet: rows with page id, image index, offers).
+"""
+import argparse, base64, glob, gzip, hashlib, io, json, math, os, random, sys
+from multiprocessing import Pool
+from pathlib import Path
+
+import numpy as np
+from PIL import Image, ImageDraw, ImageFont
+
+HUE180 = np.array([[-0.574, 1.43, 0.144], [0.426, 0.43, 0.144], [0.426, 1.43, -0.856]])
+SMALL, LARGE = "small", "large"
+GEOM = {SMALL: dict(cell=(124, 112), ctx=(232, 112), rows=8), LARGE: dict(cell=(250, 140), ctx=(250, 140), rows=5)}
+KIND = {"img": "image", "bg": "CSS background", "pseudo": "CSS icon", "svg": "inline SVG", "svgimage": "SVG image"}
+
+
+# ---------- colour maths (same as the runtime) ----------
+def lin(u):
+    u = u / 255.0
+    return np.where(u <= 0.04045, u / 12.92, ((u + 0.055) / 1.055) ** 2.4)
+
+
+def oklab(rgb):  # (..., 3) uint8 -> (..., 3)
+    r, g, b = lin(rgb[..., 0]), lin(rgb[..., 1]), lin(rgb[..., 2])
+    l = np.cbrt(0.4122214708 * r + 0.5363325363 * g + 0.0514459929 * b)
+    m = np.cbrt(0.2119034982 * r + 0.6806995451 * g + 0.1073969566 * b)
+    s = np.cbrt(0.0883024619 * r + 0.2817188376 * g + 0.6299787005 * b)
+    return np.stack([0.2104542553 * l + 0.7936177850 * m - 0.0040720468 * s,
+                     1.9779984951 * l - 2.4285922050 * m + 0.4505937099 * s,
+                     0.0259040371 * l + 0.7827717662 * m - 0.8086757660 * s], -1)
+
+
+def to_rgb(L, a, b):
+    l_ = L + 0.3963377774 * a + 0.2158037573 * b
+    m_ = L - 0.1055613458 * a - 0.0638541728 * b
+    s_ = L - 0.0894841775 * a - 1.2914855480 * b
+    l, m, s = l_ ** 3, m_ ** 3, s_ ** 3
+    def enc(x):
+        x = np.clip(x, 0, 1)
+        return np.round(255 * np.where(x <= 0.0031308, 12.92 * x, 1.055 * np.power(x, 1 / 2.4) - 0.055))
+    return np.stack([enc(4.0767416621 * l - 3.3077115913 * m + 0.2309699292 * s),
+                     enc(-1.2684380046 * l + 2.6097574011 * m - 0.3413193965 * s),
+                     enc(-0.0041960863 * l - 0.7034186147 * m + 1.7076147010 * s)], -1)
+
+
+def invert(rgba):
+    x = 1.0 - rgba[..., :3] / 255.0
+    out = rgba.copy()
+    out[..., :3] = np.round(np.clip(x @ HUE180.T, 0, 1) * 255)
+    return out
+
+
+def adapt(rgba):
+    """Port of nadir-map.js adaptPixels (incl. transparent mode). None when there is nothing to adapt."""
+    d = rgba.astype(np.float64)
+    h, w = d.shape[:2]
+    ring = np.zeros((h, w), bool)
+    ring[0, :] = ring[-1, :] = ring[:, 0] = ring[:, -1] = True
+    rp = d[ring]
+    ringN = len(rp)
+    opq = rp[:, 3] >= 250
+    n = int(opq.sum())
+    o = oklab(rp[opq][:, :3]) if n else np.zeros((0, 3))
+    lightm = (o[:, 0] > 0.9) & (np.hypot(o[:, 1], o[:, 2]) < 0.04) if n else np.zeros(0, bool)
+    light = int(lightm.sum())
+    bg = None
+    if n and light / n >= 0.5:
+        bg = oklab(rp[opq][lightm][:, :3].mean(0))
+    elif n > 0.2 * ringN:
+        return None
+    else:
+        a = d[..., 3] >= 128
+        if not a.any():
+            return None
+        ol = oklab(d[a][:, :3])
+        ink = (ol[:, 0] < 0.5) & (np.hypot(ol[:, 1], ol[:, 2]) < 0.05)
+        if ink.mean() < 0.08:
+            return None
+    out = np.zeros_like(d)
+    vis = d[..., 3] > 0
+    px = d[vis]
+    ol = oklab(px[:, :3])
+    if bg is not None:
+        dist = np.linalg.norm(ol - bg, axis=1)
+        cov = np.clip((dist - 0.03) / 0.22, 0, 1)
+    else:
+        cov = np.ones(len(px))
+    keep = cov > 0
+    c = np.where(keep, cov, 1.0)
+    if bg is not None:
+        L = (ol[:, 0] - (1 - c) * bg[0]) / c
+        A = (ol[:, 1] - (1 - c) * bg[1]) / c
+        B = (ol[:, 2] - (1 - c) * bg[2]) / c
+    else:
+        L, A, B = ol[:, 0].copy(), ol[:, 1].copy(), ol[:, 2].copy()
+    L = np.clip(L, 0, 1)
+    C = np.hypot(A, B)
+    neutral = C < 0.05
+    lift = (~neutral) & (L < 0.62)
+    L = np.where(neutral, 0.94 - 0.75 * L, np.where(lift, 0.62 + 0.15 * L / 0.62, L))
+    A = np.where(lift, A * 0.92, A)
+    B = np.where(lift, B * 0.92, B)
+    rgb = to_rgb(L, A, B)
+    res = np.zeros((len(px), 4))
+    res[:, :3] = rgb
+    res[:, 3] = np.round(px[:, 3] * cov)
+    res[~keep] = 0
+    out[vis] = res
+    return out.astype(np.uint8)
+
+
+# ---------- page data ----------
+def hex_rgba(h):
+    h = h.lstrip("#")
+    r, g, b = int(h[0:2], 16), int(h[2:4], 16), int(h[4:6], 16)
+    a = int(h[6:8], 16) / 255 if len(h) >= 8 else 1.0
+    return r, g, b, a
+
+
+def over(top, base):
+    r, g, b, a = top
+    return tuple(round(a * t + (1 - a) * u) for t, u in zip((r, g, b), base))
+
+
+def dark_palette(page, design):
+    """Dark surface per token index and the main dark text colour, from the Haiku design."""
+    cols = {int(k): v for k, v in design["colors"].items()} if isinstance(design["colors"], dict) else {int(c["id"].lstrip("t")): c["dark"] for c in design["colors"]}
+    toks = page["tokens"]
+    root = next((i for i, t in enumerate(toks) if t["role"] == 0 and t["f"][24] > 0.5), 0)
+    canvas = hex_rgba(cols[root])[:3] if root in cols else (18, 18, 20)
+    best, text = -1, (230, 232, 235)
+    for i, t in enumerate(toks):
+        if t["role"] == 1 and i in cols and t["f"][19] > best:
+            best, text = t["f"][19], over(hex_rgba(cols[i]), canvas)
+    surf = {i: over(hex_rgba(cols[i]), canvas) for i, t in enumerate(toks) if i in cols and t["role"] == 0}
+    return surf, canvas, text
+
+
+def content_area(o, big):
+    """The big thumbnail is letterboxed into a square: cut out the image itself."""
+    S = big.shape[0]
+    if o.get("cropPx"):
+        nw, nh = o["cropPx"][2], o["cropPx"][3]
+    elif o.get("kind") == "svg":
+        nw, nh = o["w"], o["h"]
+    else:
+        nw, nh = (o.get("nw") or o["w"]), (o.get("nh") or o["h"])
+    nw, nh = max(1, nw), max(1, nh)
+    s = min(S / nw, S / nh)
+    dw, dh = max(1, round(nw * s)), max(1, round(nh * s))
+    x0, y0 = (S - dw) // 2, (S - dh) // 2
+    return big[y0:y0 + dh, x0:x0 + dw]
+
+
+def decode_png(data_url):
+    return np.array(Image.open(io.BytesIO(base64.b64decode(data_url.split(",", 1)[1]))).convert("RGBA"))
+
+
+def bucket(o):
+    s = o.get("stats") or {}
+    area = o["w"] * o["h"]
+    if o.get("logo") or (o.get("homeLink") and o.get("header")):
+        return "logo"
+    if o["kind"] == "bg" and area >= 20000:
+        return "backdrop"
+    if o["kind"] in ("pseudo",) or o.get("cropPx"):
+        return "icon"
+    if s.get("transp", 0) > 0.15:
+        return "transparent"
+    if s.get("transp", 1) < 0.02 and s.get("colors", 0) > 0.18 and s.get("stdL", 0) > 0.12:
+        return "photo"
+    return "flat"
+
+
+RATE = {"logo": 1.0, "backdrop": 1.0, "icon": 1.0, "transparent": 1.0, "flat": 0.8, "photo": 0.12}
+
+
+def candidates(task):
+    """All sampled elements of one page with everything needed to draw their rows."""
+    crawl, haiku, pid, tag, seed = task
+    rnd = random.Random(seed)
+    try:
+        page = json.load(gzip.open(Path(crawl) / "pages" / pid / "light.json.gz"))
+        design = json.load(open(Path(haiku) / f"{pid}.json"))
+        if design.get("refusal") or not design.get("colors"):
+            return []
+        surf, canvas, text = dark_palette(page, design)
+    except Exception:
+        return []
+    out = []
+    for i, o in enumerate(page["images"]):
+        if not o.get("big") or o["w"] * o["h"] < 64 or o["w"] < 4 or o["h"] < 4:
+            continue
+        if o["kind"] == "svg" and not o.get("useExt"):  # inline SVGs are recoloured by the colour model, not as images
+            continue
+        b = bucket(o)
+        if rnd.random() > RATE[b]:
+            continue
+        dsurf = surf.get(o.get("underTok", -1), canvas)
+        out.append({"page": pid, "crawl": tag, "img": i, "bucket": b, "o": {k: o.get(k) for k in ("kind", "src", "x", "y", "w", "h", "nw", "nh", "cropPx", "logo", "homeLink", "header", "inLink", "surf", "big", "stats")},
+                    "dsurf": dsurf, "dtext": text, "light": tuple((o.get("surf") or [255, 255, 255])[:3])})
+    return out
+
+
+# ---------- drawing ----------
+def font(sz):
+    try:
+        return ImageFont.load_default(size=sz)
+    except TypeError:
+        return ImageFont.load_default()
+
+
+F11, F13, F15 = font(11), font(13), font(15)
+
+
+def fit(o, box, maxzoom=4.0):
+    s = min(box[0] / max(1, o["w"]), box[1] / max(1, o["h"]), maxzoom)
+    return max(1, round(o["w"] * s)), max(1, round(o["h"] * s)), s
+
+
+def variant_cell(img, o, typ, size, surface, text, mode):
+    W, H = size
+    cell = Image.new("RGBA", (W, H), tuple(surface) + (255,))
+    if mode != "D":
+        dw, dh, _ = fit(o, (W - 12, H - 12))
+        im = Image.fromarray(img, "RGBA").resize((dw, dh), Image.LANCZOS)
+        cell.alpha_composite(im, ((W - dw) // 2, (H - dh) // 2))
+    if typ == LARGE and o["kind"] == "bg":  # content on a background image: in dark mode its text is light
+        dr = ImageDraw.Draw(cell)
+        dr.text((12, H // 2 - 22), "Heading on this background", fill=tuple(text), font=F15)
+        dr.text((12, H // 2 + 2), "Body text that sits on top of it.", fill=tuple(text), font=F13)
+    return cell
+
+
+def context_cell(shot, e, size):
+    W, H = size
+    o = e["o"]
+    if shot is not None and o["y"] + o["h"] <= shot.height and o["x"] >= -2 and o["x"] + o["w"] <= shot.width + 2:
+        m = max(30, 0.75 * max(o["w"], o["h"]))
+        x0, y0 = max(0, o["x"] - m), max(0, o["y"] - m)
+        x1, y1 = min(shot.width, o["x"] + o["w"] + m), min(shot.height, o["y"] + o["h"] + m)
+        crop = shot.crop((int(x0), int(y0), int(x1), int(y1)))
+        s = min(W / crop.width, H / crop.height, 3.0)
+        cw, ch = max(1, round(crop.width * s)), max(1, round(crop.height * s))
+        crop = crop.resize((cw, ch), Image.LANCZOS).convert("RGBA")
+        cell = Image.new("RGBA", (W, H), (236, 236, 236, 255))
+        ox, oy = (W - cw) // 2, (H - ch) // 2
+        cell.alpha_composite(crop, (ox, oy))
+        dr = ImageDraw.Draw(cell)
+        rx, ry = ox + (o["x"] - x0) * s, oy + (o["y"] - y0) * s
+        dr.rectangle([rx - 2, ry - 2, rx + o["w"] * s + 1, ry + o["h"] * s + 1], outline=(230, 20, 20, 255), width=2)
+        return cell
+    # below the captured screenshot: the element on its light surface
+    cell = Image.new("RGBA", (W, H), tuple(e["light"]) + (255,))
+    big = content_area(o, decode_png(o["big"]))
+    dw, dh, _ = fit(o, (W - 16, H - 16))
+    cell.alpha_composite(Image.fromarray(big, "RGBA").resize((dw, dh), Image.LANCZOS), ((W - dw) // 2, (H - dh) // 2))
+    ImageDraw.Draw(cell).text((4, H - 14), "(on its light background)", fill=(90, 90, 90), font=F11)
+    return cell
+
+
+def draw_sheet(rows, typ, shots, path):
+    g = GEOM[typ]
+    cw, ch = g["cell"]
+    xw, xh = g["ctx"]
+    pad, lab, cap = 6, 30, 16
+    W = lab + xw + 4 * (cw + pad) + pad
+    rh = max(ch, xh) + cap + 10
+    H = 26 + rh * len(rows)
+    sheet = Image.new("RGB", (W, H), (255, 255, 255))
+    dr = ImageDraw.Draw(sheet)
+    heads = ["light page (red box)", "A keep", "B invert", "C adapt", "D remove"]
+    xs = [lab, lab + xw + pad] + [lab + xw + pad + k * (cw + pad) for k in range(1, 4)]
+    for x, t in zip(xs, heads):
+        dr.text((x + 4, 6), t, fill=(0, 0, 0), font=F13)
+    meta = []
+    for r, e in enumerate(rows):
+        y = 26 + r * rh
+        o = e["o"]
+        dr.text((4, y + ch // 2 - 8), f"#{r + 1}", fill=(0, 0, 0), font=F15)
+        sheet.paste(context_cell(shots.get(e["page"]), e, (xw, xh)).convert("RGB"), (xs[0], y))
+        big = content_area(o, decode_png(o["big"]))
+        offers = "AB"
+        inv = invert(big)
+        ad = adapt(big)
+        cells = [variant_cell(big, o, typ, (cw, ch), e["dsurf"], e["dtext"], "A"), variant_cell(inv, o, typ, (cw, ch), e["dsurf"], e["dtext"], "B")]
+        if ad is not None:
+            offers += "C"
+            cells.append(variant_cell(ad, o, typ, (cw, ch), e["dsurf"], e["dtext"], "C"))
+        else:
+            cells.append(None)
+        if typ == LARGE and o["kind"] == "bg":
+            offers += "D"
+            cells.append(variant_cell(big, o, typ, (cw, ch), e["dsurf"], e["dtext"], "D"))
+        else:
+            cells.append(None)
+        for k, c in enumerate(cells):
+            x = xs[1 + k]
+            if c is None:
+                dr.rectangle([x, y, x + cw - 1, y + ch - 1], fill=(225, 225, 225))
+                dr.text((x + cw // 2 - 14, y + ch // 2 - 7), "n/a", fill=(140, 140, 140), font=F13)
+            else:
+                sheet.paste(c.convert("RGB"), (x, y))
+        flags = [f for f, v in (("logo", o.get("logo")), ("links home", o.get("homeLink")), ("in a link", o.get("inLink") and not o.get("homeLink")), ("in header", o.get("header"))) if v]
+        _, _, z = fit(o, (cw - 12, ch - 12))
+        zoom = f", shown {z:.1f}x" if abs(z - 1) > 0.15 else ""
+        dr.text((xs[1], y + ch + 2), f"#{r + 1}: {KIND.get(o['kind'], o['kind'])} {o['w']}x{o['h']} px{zoom}" + (" · " + ", ".join(flags) if flags else ""), fill=(60, 60, 60), font=F11)
+        meta.append({"k": r + 1, "page": e["page"], "crawl": e["crawl"], "img": e["img"], "bucket": e["bucket"], "offer": offers, "kind": o["kind"]})
+    sheet.save(path, optimize=True)
+    return meta
+
+
+def render_chunk(job):
+    wid, rows_by_sheet, crawls, out = job
+    shots = {}
+    outp = Path(out)
+    lines = []
+    for sid, typ, rows in rows_by_sheet:
+        for e in rows:
+            if e["page"] not in shots:
+                p = Path(crawls[e["crawl"]]) / "pages" / e["page"] / "shot.jpg"
+                try:
+                    shots[e["page"]] = Image.open(p).convert("RGB")
+                except Exception:
+                    shots[e["page"]] = None
+        try:
+            meta = draw_sheet(rows, typ, shots, outp / "sheets" / f"{sid}.png")
+            lines.append(json.dumps({"sheet": sid, "type": typ, "rows": meta}))
+        except Exception as ex:
+            print("sheet error", sid, repr(ex)[:200], flush=True)
+        if len(shots) > 64:
+            shots.clear()
+    with open(outp / f"meta.{wid}.jsonl", "a") as f:
+        f.write("\n".join(lines) + ("\n" if lines else ""))
+    return len(lines)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--crawl", action="append", required=True, help="TAG=CRAWL_DIR=HAIKU_DIR (repeatable)")
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--max-elements", type=int, default=150000)
+    ap.add_argument("--max-per-src", type=int, default=2)
+    ap.add_argument("--workers", type=int, default=28)
+    ap.add_argument("--exclude-pages", default="", help="file of page ids already done")
+    ap.add_argument("--flat-rate", type=float, default=None, help="sampling rate for opaque flat/photo-like images (mostly 'keep')")
+    args = ap.parse_args()
+    if args.flat_rate is not None:
+        RATE["flat"] = args.flat_rate
+    out = Path(args.out)
+    (out / "sheets").mkdir(parents=True, exist_ok=True)
+    crawls, tasks = {}, []
+    done = set(l.strip() for l in open(args.exclude_pages)) if args.exclude_pages and os.path.exists(args.exclude_pages) else set()
+    for spec in args.crawl:
+        tag, cdir, hdir = spec.split("=")
+        crawls[tag] = cdir
+        for hp in Path(hdir).glob("*.json"):
+            pid = hp.stem
+            if len(pid) == 16 and pid not in done and (Path(cdir) / "pages" / pid / "light.json.gz").exists():
+                tasks.append((cdir, hdir, pid, tag, int(hashlib.md5(pid.encode()).hexdigest()[:8], 16)))
+    random.Random(3).shuffle(tasks)
+    print(len(tasks), "pages with designs", flush=True)
+    per_src, elems, by_bucket = {}, [], {}
+    with Pool(args.workers) as pool:
+        for res in pool.imap_unordered(candidates, tasks, chunksize=16):
+            for e in res:
+                s = e["o"]["src"]
+                if per_src.get(s, 0) >= args.max_per_src:
+                    continue
+                per_src[s] = per_src.get(s, 0) + 1
+                elems.append(e)
+                by_bucket[e["bucket"]] = by_bucket.get(e["bucket"], 0) + 1
+            if len(elems) >= args.max_elements:
+                break
+    elems = elems[: args.max_elements]
+    print(len(elems), "elements", by_bucket, flush=True)
+    with open(out / "pages_used.txt", "a") as f:
+        f.write("\n".join(sorted(set(e["page"] for e in elems))) + "\n")
+    # sheets: small and large elements apart; rows mixed across pages
+    random.Random(5).shuffle(elems)
+    groups = {SMALL: [], LARGE: []}
+    for e in elems:
+        o = e["o"]
+        typ = LARGE if (o["w"] * o["h"] >= 20000 or o["w"] > 300) else SMALL
+        groups[typ].append(e)
+    sheets = []
+    for typ, es in groups.items():
+        n = GEOM[typ]["rows"]
+        for i in range(0, len(es), n):
+            rows = es[i:i + n]
+            sid = hashlib.sha1(json.dumps([(r["page"], r["img"]) for r in rows]).encode()).hexdigest()[:16]
+            sheets.append((sid, typ, rows))
+    jobs = [(w, sheets[w::args.workers], crawls, str(out)) for w in range(args.workers)]
+    with Pool(args.workers) as pool:
+        n = sum(pool.map(render_chunk, jobs))
+    print("sheets", n, flush=True)
+
+
+if __name__ == "__main__":
+    main()

training/export.py

@@ -1,15 +1,18 @@
 #!/usr/bin/env python3
-"""Export NadirNet to ONNX (fp32/fp16) and build calibrated 8-bit and 4-bit variants; measure each against fp32."""
+"""Export NadirNet to ONNX: the page graph (palette + page context) and the element graph (treatment + role for any
+set of page images). Calibrated 8-bit and 4-bit variants of the page graph come from quantize.py."""
 import argparse, json, os, pickle, time
 import numpy as np
 import torch
 import torch.nn as nn
 
-from model import NadirNet
-from train import collate, prep, evaluate
+from model import NadirNet, ElemNet, FI, TREAT, IMG_ROLES, N_CTX
+from train import collate, prep, element_logits
 
 NAMES_IN = ["tok_f", "tok_role", "tok_lab", "img_f", "img_thumb", "page_f", "src", "rel", "mask"]
-NAMES_OUT = ["lab", "alpha_logit", "img_logits"]
+NAMES_OUT = ["lab", "alpha_logit", "ctx"]
+ELEM_IN = ["e_f", "e_thumb", "under_light", "under_dark", "ctx"]
+ELEM_OUT = ["treat_logits", "role_logits"]
 
 
 class Wrapped(nn.Module):
@@ -45,16 +48,25 @@ def feeds_for(s, pad=None):
     M = b["img_f"].shape[1]
     if M == 0:  # the runtime always sends at least one (masked) image slot
         T = b["mask"].shape[1]
-        b["img_f"] = torch.zeros(1, 1, 39); b["thumbs"] = torch.zeros(1, 1, 4, 32, 32)
+        b["img_f"] = torch.zeros(1, 1, FI); b["thumbs"] = torch.zeros(1, 1, 4, 32, 32)
         rel = torch.zeros(1, T + 1, T + 1, b["rel"].shape[-1]); rel[:, :T, :T] = b["rel"]; b["rel"] = rel
         b["mask"] = torch.cat([b["mask"], torch.zeros(1, 1, dtype=torch.bool)], 1)
     return {"tok_f": b["tok_f"].numpy(), "tok_role": b["role"].numpy(), "tok_lab": b["lab"].numpy(), "img_f": b["img_f"].numpy(),
             "img_thumb": b["thumbs"].numpy(), "page_f": b["page"].numpy(), "src": b["src"].numpy(), "rel": b["rel"].numpy(), "mask": b["mask"].numpy()}
 
 
+def elem_feeds(s, lab, ctx):
+    """Element-graph inputs for a page's labelled images, with the dark surface taken from the page graph's output."""
+    b = collate([s], "cpu")
+    k = int(b["e_mask"][0].sum())
+    ut = b["e_ut"][0, :k].numpy()
+    return {"e_f": b["e_f"][:, :k].numpy(), "e_thumb": b["e_thumb"][:, :k].numpy(), "under_light": b["e_ul"][:, :k].numpy(),
+            "under_dark": lab[:, ut].astype(np.float32), "ctx": ctx.astype(np.float32)}, b["e_soft"][0, :k].numpy()
+
+
 def fidelity(sess, ref_sess, samples):
-    """Mean OKLab distance between this model's colours and fp32's, plus image-decision agreement."""
-    de, agree, n_img, worst = [], 0, 0, 0.0
+    """Mean OKLab distance between this model's colours and fp32's (weighted like training), worst token, ctx drift."""
+    de, worst, ctxd = [], 0.0, []
     for s in samples:
         f = feeds_for(s)
         a = sess.run(None, f); r = ref_sess.run(None, f)
@@ -62,16 +74,13 @@ def fidelity(sess, ref_sess, samples):
         d = np.linalg.norm(a[0][0, :n] - r[0][0, :n], axis=-1)
         w = s["w"][:n]
         de.append(float((d * w).sum() / w.sum())); worst = max(worst, float(d.max()))
-        m = len(s["IF"])
-        if m:
-            agree += int((a[2][0, :m].argmax(-1) == r[2][0, :m].argmax(-1)).sum()); n_img += m
-    return {"dE_vs_fp32": float(np.mean(de)), "dE_max": worst, "img_agree": agree / max(1, n_img)}
+        ctxd.append(float(np.abs(a[2] - r[2]).max()))
+    return {"dE_vs_fp32": float(np.mean(de)), "dE_max": worst, "ctx_max_abs": float(np.max(ctxd))}
 
 
-def bench(sess, samples, reps=3):
+def bench(sess, feeds, reps=3):
     ts = []
-    for s in samples[:20]:
-        f = feeds_for(s)
+    for f in feeds[:20]:
         sess.run(None, f)
         t = time.time()
         for _ in range(reps):
@@ -85,8 +94,7 @@ def main():
     ap.add_argument("--ckpt", required=True)
     ap.add_argument("--data", required=True)
     ap.add_argument("--out", required=True)
-    ap.add_argument("--calib", type=int, default=256)
-    ap.add_argument("--only-float", action="store_true", help="fp32/fp16 only (calibrated 8/4-bit come from quantize.py)")
+    ap.add_argument("--only-float", action="store_true", help="kept for the job scripts (quantization lives in quantize.py)")
     args = ap.parse_args()
     os.makedirs(args.out, exist_ok=True)
     import onnx, onnxruntime as ort
@@ -94,116 +102,56 @@ def main():
     model = NadirNet(**{k: v for k, v in ck["cfg"].items()}); model.load_state_dict(ck["model"]); model.eval()
     data = pickle.load(open(args.data, "rb"))
     val = [prep(s) for s in data["val"]]
-    rng = np.random.default_rng(0)
-    calib = [prep(s) for s in rng.choice(np.array(data["train"], dtype=object), size=min(args.calib, len(data["train"])), replace=False)]
     s0 = max(val, key=lambda s: len(s["role"]))
     ex = feeds_for(s0)
     par_samples = val[:8]
     with torch.no_grad():
-        ref_out = [model(*[torch.from_numpy(feeds_for(s)[k]) for k in NAMES_IN])[0].numpy() for s in par_samples]
+        ref_out = [model(*[torch.from_numpy(feeds_for(s)[k]) for k in NAMES_IN]) for s in par_samples]
     fp32 = os.path.join(args.out, "nadir-fp32.onnx")
     torch.onnx.export(Wrapped(model), tuple(torch.from_numpy(ex[k]) for k in NAMES_IN), fp32, input_names=NAMES_IN, output_names=NAMES_OUT,
                       dynamic_axes={"tok_f": {1: "N"}, "tok_role": {1: "N"}, "tok_lab": {1: "N"}, "img_f": {1: "M"}, "img_thumb": {1: "M"},
-                                    "rel": {1: "T", 2: "T"}, "mask": {1: "T"}, "lab": {1: "N"}, "alpha_logit": {1: "N"}, "img_logits": {1: "M"}},
+                                    "rel": {1: "T", 2: "T"}, "mask": {1: "T"}, "lab": {1: "N"}, "alpha_logit": {1: "N"}},
                       opset_version=17, do_constant_folding=True, dynamo=False)
     onnx.checker.check_model(fp32)
     so = ort.SessionOptions(); so.intra_op_num_threads = 4
     ref = ort.InferenceSession(fp32, so, providers=["CPUExecutionProvider"])
-    # parity torch vs onnx (reference computed before export)
-    parity = max(float(np.abs(r - ref.run(None, feeds_for(s))[0]).max()) for r, s in zip(ref_out, par_samples))
-    report = {"params": sum(p.numel() for p in model.parameters()), "parity_max_abs": parity, "variants": {}}
-    print("parity", parity, flush=True)
-
-    # fp16 (weights + compute in fp16, IO kept fp32)
-    # fp16: export a half-precision copy with fp32 inputs/outputs (casts inside the graph)
-    have16 = True
-    try:
-        class Half(nn.Module):
-            def __init__(self, m):
-                super().__init__(); self.m = m.half()
-            def forward(self, tok_f, tok_role, tok_lab, img_f, img_thumb, page_f, src, rel, mask):
-                o = self.m(tok_f.half(), tok_role, tok_lab.half(), img_f.half(), img_thumb.half(), page_f.half(), src, rel.half(), mask)
-                return tuple(x.float() for x in o)
-        import copy
-        torch.onnx.export(Half(copy.deepcopy(model)), tuple(torch.from_numpy(ex[k]) for k in NAMES_IN), os.path.join(args.out, "nadir-fp16.onnx"),
-                          input_names=NAMES_IN, output_names=NAMES_OUT,
-                          dynamic_axes={"tok_f": {1: "N"}, "tok_role": {1: "N"}, "tok_lab": {1: "N"}, "img_f": {1: "M"}, "img_thumb": {1: "M"},
-                                        "rel": {1: "T", 2: "T"}, "mask": {1: "T"}, "lab": {1: "N"}, "alpha_logit": {1: "N"}, "img_logits": {1: "M"}},
-                          opset_version=17, do_constant_folding=True, dynamo=False)
-        ort.InferenceSession(os.path.join(args.out, "nadir-fp16.onnx"), providers=["CPUExecutionProvider"])
-    except Exception as e:
-        print("fp16 failed:", str(e)[:300], flush=True); have16 = False
-
-    if args.only_float:
-        variants = {"fp32": fp32}
-        if have16: variants["fp16"] = os.path.join(args.out, "nadir-fp16.onnx")
-        for name, path in variants.items():
-            sess = ort.InferenceSession(path, so, providers=["CPUExecutionProvider"])
-            r = fidelity(sess, ref, val[:300]); r["size_mb"] = round(os.path.getsize(path) / 1e6, 2); r["cpu_ms_median"] = round(bench(sess, val[:300]), 2)
-            report["variants"][name] = r
-            print(name, r, flush=True)
-        json.dump(report, open(os.path.join(args.out, "export_report.json"), "w"), indent=1)
-        return
-    from onnxruntime.quantization import quantize_dynamic, quantize_static, QuantType, QuantFormat, CalibrationDataReader, CalibrationMethod
-    from onnxruntime.quantization.shape_inference import quant_pre_process
-    pre = os.path.join(args.out, "_pre.onnx")
-    quant_pre_process(fp32, pre, skip_symbolic_shape=True)
-
-    class Reader(CalibrationDataReader):
-        def __init__(self, ss):
-            self.it = iter([feeds_for(s, pad=(256, 32)) for s in ss])
-        def get_next(self):
-            return next(self.it, None)
-
-    # 8-bit A: per-channel int8 weights, dynamic activations
-    q8d = os.path.join(args.out, "nadir-int8-dynamic.onnx")
-    quantize_dynamic(pre, q8d, weight_type=QuantType.QInt8, per_channel=True, op_types_to_quantize=["MatMul", "Gemm"])
-    # 8-bit B: static QDQ, activations calibrated on real pages (percentile)
-    q8s = os.path.join(args.out, "nadir-int8-static.onnx")
-    try:
-        quantize_static(pre, q8s, Reader(calib[:128]), quant_format=QuantFormat.QDQ, per_channel=True, activation_type=QuantType.QInt8,
-                        weight_type=QuantType.QInt8, op_types_to_quantize=["MatMul", "Gemm"], calibrate_method=CalibrationMethod.Percentile,
-                        extra_options={"CalibPercentile": 99.99, "ActivationSymmetric": True, "WeightSymmetric": True})
-    except Exception as e:
-        print("static int8 failed:", e, flush=True); q8s = None
-    # 4-bit: block-wise weights (MatMulNBits), GPTQ calibrated on real pages, RTN fallback
-    q4 = os.path.join(args.out, "nadir-int4.onnx")
-    from onnxruntime.quantization import matmul_nbits_quantizer as mnq
-    done4 = False
-    # GPTQ needs reduction dims that are multiples of the block size; leave tiny matrices (relation bias, 11 wide) in fp32
-    pm = onnx.load(pre)
-    inits = {i.name: i for i in pm.graph.initializer}
-    skip = [n.name for n in pm.graph.node if n.op_type == "MatMul" and len(n.input) > 1 and n.input[1] in inits and inits[n.input[1]].dims[0] % 32 != 0]
-    report["int4_fp32_nodes"] = len(skip)
-    for algo in ("gptq", "hqq", "rtn"):
-        try:
-            if algo == "gptq":
-                cfg = mnq.GPTQWeightOnlyQuantConfig(calibration_data_reader=Reader(calib), percdamp=0.01, block_size=32, actorder=True, mse=True)
-            elif algo == "hqq":
-                cfg = mnq.HQQWeightOnlyQuantConfig(block_size=32, bits=4)
-            else:
-                cfg = mnq.RTNWeightOnlyQuantConfig()
-            qz = mnq.MatMulNBitsQuantizer(onnx.load(pre), block_size=32, is_symmetric=False, accuracy_level=4, algo_config=cfg, nodes_to_exclude=skip)
-            qz.process()
-            qz.model.save_model_to_file(q4, use_external_data_format=False)
-            report["int4_algo"] = algo
-            done4 = True
-            break
-        except Exception as e:
-            print(f"int4 {algo} failed: {str(e)[:300]}", flush=True)
-    variants = {"fp32": fp32, "int8-dynamic": q8d}
-    if have16: variants["fp16"] = os.path.join(args.out, "nadir-fp16.onnx")
-    if q8s: variants["int8-static"] = q8s
-    if done4: variants["int4"] = q4
-    evs = val[:300]
-    for name, path in variants.items():
-        sess = ort.InferenceSession(path, so, providers=["CPUExecutionProvider"])
-        r = fidelity(sess, ref, evs)
-        r["size_mb"] = round(os.path.getsize(path) / 1e6, 2)
-        r["cpu_ms_median"] = round(bench(sess, evs), 2)
-        report["variants"][name] = r
-        print(name, r, flush=True)
-    os.remove(pre)
+    parity = 0.0
+    for r, s in zip(ref_out, par_samples):
+        o = ref.run(None, feeds_for(s))
+        parity = max(parity, float(np.abs(r[0].numpy() - o[0]).max()), float(np.abs(r[2].numpy() - o[2]).max()))
+    # element graph
+    es = [s for s in val if len(s["E_ut"])][:64]
+    ef = []
+    for s in es:
+        o = ref.run(None, feeds_for(s))
+        ef.append(elem_feeds(s, o[0], o[2]))
+    exe = max(ef, key=lambda x: x[0]["e_f"].shape[1])[0] if ef else {"e_f": np.zeros((1, 4, FI), np.float32), "e_thumb": np.zeros((1, 4, 4, 32, 32), np.float32),
+                                                                    "under_light": np.zeros((1, 4, 3), np.float32), "under_dark": np.zeros((1, 4, 3), np.float32), "ctx": np.zeros((1, N_CTX), np.float32)}
+    elem = os.path.join(args.out, "nadir-elem-fp32.onnx")
+    em = ElemNet(model).eval()
+    torch.onnx.export(em, tuple(torch.from_numpy(exe[k]) for k in ELEM_IN), elem, input_names=ELEM_IN, output_names=ELEM_OUT,
+                      dynamic_axes={"e_f": {1: "K"}, "e_thumb": {1: "K"}, "under_light": {1: "K"}, "under_dark": {1: "K"}, "treat_logits": {1: "K"}, "role_logits": {1: "K"}},
+                      opset_version=17, do_constant_folding=True, dynamo=False)
+    onnx.checker.check_model(elem)
+    esess = ort.InferenceSession(elem, so, providers=["CPUExecutionProvider"])
+    eparity, agree, ok, n = 0.0, 0, 0, 0
+    T = float(ck.get("temperature", 1.0))
+    for f, soft in ef:
+        with torch.no_grad():
+            tt, rr = em(*[torch.from_numpy(f[k]) for k in ELEM_IN])
+        o = esess.run(None, f)
+        eparity = max(eparity, float(np.abs(tt.numpy() - o[0]).max()), float(np.abs(rr.numpy() - o[1]).max()))
+        pred = o[0][0].argmax(-1)
+        agree += int((pred == soft.argmax(-1)).sum()); ok += int((soft[np.arange(len(pred)), pred] > 0).sum()); n += len(pred)
+    report = {"params": sum(p.numel() for p in model.parameters()),
+              "elem_params": sum(p.numel() for k, p in model.named_parameters() if k.startswith(("cnn.", "img_in.", "elem_head."))),
+              "parity_max_abs": parity, "elem_parity_max_abs": eparity, "temperature": T, "treat": TREAT, "img_roles": IMG_ROLES,
+              "elem_val_sample": {"n": n, "acc": agree / max(1, n), "ok": ok / max(1, n)}, "variants": {}}
+    print("parity", parity, "elem parity", eparity, flush=True)
+    r = fidelity(ref, ref, val[:50]); r["size_mb"] = round(os.path.getsize(fp32) / 1e6, 2); r["cpu_ms_median"] = round(bench(ref, [feeds_for(s) for s in val[:20]]), 2)
+    report["variants"]["fp32"] = r
+    report["elem"] = {"size_mb": round(os.path.getsize(elem) / 1e6, 2), "cpu_ms_median": round(bench(esess, [f for f, _ in ef]), 2) if ef else None}
+    print(report, flush=True)
     json.dump(report, open(os.path.join(args.out, "export_report.json"), "w"), indent=1)
 
 

training/ext_render.py

@@ -0,0 +1,87 @@
+#!/usr/bin/env python3
+"""End-to-end eval renders: each page as the user sees it with a given build of the extension (or none).
+
+usage: ext_render.py --urls URLS.jsonl --out OUT --variant TAG=EXT_DIR [--variant TAG=EXT_DIR ...] [--light]
+URLS.jsonl lines: {"id", "url"}. Writes OUT/<tag>/<id>.jpg (viewport, after the extension settled) and
+OUT/<tag>/<id>.json (Nadir status). With --light also OUT/light/<id>.jpg without any extension.
+Each variant runs in its own persistent profile, so first-visit behaviour (no site cache) is what gets judged.
+"""
+import argparse, asyncio, json, os, shutil, time
+from pathlib import Path
+from playwright.async_api import async_playwright
+
+STATUS_JS = "async () => { const a = await chrome.storage.session.get(null); return Object.values(a).pop() || null; }"
+
+
+async def render_variant(p, tag, ext, urls, out, args):
+    od = out / tag
+    od.mkdir(parents=True, exist_ok=True)
+    prof = out / f"_profile_{tag}"
+    shutil.rmtree(prof, ignore_errors=True)
+    # the whole browser (extension pages too) presents a normal Chrome UA: some image hosts refuse "HeadlessChrome"
+    launch_args = ["--hide-scrollbars", "--mute-audio", "--disable-features=DisableLoadExtensionCommandLineSwitch", f"--user-agent={args.ua}"]
+    if ext:
+        launch_args += [f"--disable-extensions-except={ext}", f"--load-extension={ext}"]
+    # extensions need the full browser in new-headless mode, not Playwright's headless shell
+    ctx = await p.chromium.launch_persistent_context(str(prof), executable_path=args.exe or None, channel=None if args.exe else "chromium", headless=True,
+                                                     ignore_default_args=["--disable-extensions"] if ext else [],
+                                                     viewport={"width": 1280, "height": 800}, args=launch_args, user_agent=args.ua)
+    sw = None
+    if ext:
+        sw = ctx.service_workers[0] if ctx.service_workers else await ctx.wait_for_event("serviceworker", timeout=20000)
+    sem = asyncio.Semaphore(args.conc)
+
+    async def one(rec):
+        dst = od / f"{rec['id']}.jpg"
+        if dst.exists():
+            return
+        async with sem:
+            pg = await ctx.new_page()
+            st = {"url": rec["url"], "t": time.time()}
+            try:
+                await pg.goto(rec["url"], wait_until="domcontentloaded", timeout=30000)
+                try:
+                    await pg.wait_for_load_state("load", timeout=10000)
+                except Exception:
+                    pass
+                await pg.wait_for_timeout(args.wait * 1000)
+                await pg.screenshot(path=str(dst), type="jpeg", quality=82)
+                if sw:
+                    try:
+                        st["nadir"] = await sw.evaluate(STATUS_JS)
+                    except Exception as e:
+                        st["nadir_err"] = str(e)[:120]
+            except Exception as e:
+                st["error"] = str(e)[:200]
+            finally:
+                st["ok"] = dst.exists() and "error" not in st
+                (od / f"{rec['id']}.json").write_text(json.dumps(st))
+                await pg.close()
+    await asyncio.gather(*[one(r) for r in urls])
+    await ctx.close()
+    shutil.rmtree(prof, ignore_errors=True)
+
+
+async def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--urls", required=True)
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--variant", action="append", default=[], help="TAG=EXT_DIR")
+    ap.add_argument("--light", action="store_true")
+    ap.add_argument("--exe", default="")
+    ap.add_argument("--wait", type=float, default=7.0)
+    ap.add_argument("--conc", type=int, default=4)
+    ap.add_argument("--ua", default="Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/141.0.0.0 Safari/537.36")
+    args = ap.parse_args()
+    out = Path(args.out)
+    urls = [json.loads(l) for l in open(args.urls) if l.strip()]
+    async with async_playwright() as p:
+        jobs = [("light", None)] if args.light else []
+        jobs += [tuple(v.split("=", 1)) for v in args.variant]
+        # variants side by side (separate browsers): the same live page state as far as possible
+        await asyncio.gather(*[render_variant(p, tag, ext, urls, out, args) for tag, ext in jobs])
+    print("done", len(urls), [t for t, _ in jobs], flush=True)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

training/haiku_gen.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 """Nadir teacher data: Claude Haiku 5.5 designs a dark theme for each crawled light page.
 
-Runs on Vega's host (the API key vault proxy is only reachable there). Writes one JSON per page to
+Reads ANTHROPIC_API_KEY from the environment. Writes one JSON per page to
 OUT/<id>.json and keeps a running cost ledger; stops at --budget dollars.
 """
 import argparse, asyncio, base64, gzip, io, json, math, os, random, re, sys, time

training/judge.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-"""Blind pairwise judging of dark-mode renders with Claude Haiku 5.5 (Batch API). Runs on Vega's host."""
+"""Blind pairwise judging of dark-mode renders with Claude Haiku 5.5 (Batch API). Reads ANTHROPIC_API_KEY via haiku_gen."""
 import argparse, base64, io, json, random, time
 from pathlib import Path
 from PIL import Image
@@ -32,6 +32,7 @@ def main():
     ap.add_argument("--effort", default="medium")
     ap.add_argument("--budget", type=float, default=3.0)
     ap.add_argument("--chunk", type=int, default=400, help="requests per batch (3 images each, ~150 KB)")
+    ap.add_argument("--light-tag", default="", help="reference light renders from RENDERS/<tag>/<id>.jpg instead of the crawl screenshot")
     args = ap.parse_args()
     R, crawl, out = Path(args.renders), Path(args.crawl), Path(args.out)
     out.mkdir(parents=True, exist_ok=True)
@@ -50,7 +51,10 @@ def main():
             flip = rng.random() < 0.5
             a, b = (y, x) if flip else (x, y)
             cid = f"{pid}-{x}-{y}"
-            content = [{"type": "text", "text": "Original light page:"}, img_block(crawl / "pages" / pid / "shot.jpg"),
+            ref = (R / args.light_tag / f"{pid}.jpg") if args.light_tag else (crawl / "pages" / pid / "shot.jpg")
+            if not ref.exists():
+                continue
+            content = [{"type": "text", "text": "Original light page:"}, img_block(ref),
                        {"type": "text", "text": "Dark mode A:"}, img_block(R / a / f"{pid}.jpg"),
                        {"type": "text", "text": "Dark mode B:"}, img_block(R / b / f"{pid}.jpg"),
                        {"type": "text", "text": "Which dark mode is better? Scores are 1-10."}]

training/make_urls2.py

@@ -0,0 +1,71 @@
+#!/usr/bin/env python3
+"""Round-2 crawl list: inner pages (articles, docs, product and app pages) of the sites crawled in round 1.
+
+Round 1 crawled home pages only; the extension mostly meets inner pages. One same-site link per light or
+native-dark-mode site (a second one for some), preferring content-like paths, plus a curated list of web apps.
+"""
+import argparse, glob, json, random, re
+from urllib.parse import urlparse
+
+CONTENT = re.compile(r"/(blog|news|article|articles|post|posts|docs?|documentation|help|support|guide|guides|learn|wiki|product|products|"
+                     r"shop|store|category|categories|topic|topics|forum|community|about|pricing|features|faq|story|stories|events?|"
+                     r"courses?|recipes?|reviews?|search|tag|tags|catalog|collections?|services?)(/|$)", re.I)
+
+
+def score(u):
+    p = urlparse(u).path.rstrip("/")
+    segs = [s for s in p.split("/") if s]
+    s = 1.0 + (2.0 if CONTENT.search(p + "/") else 0) + (1.0 if len(segs) >= 2 else 0) + (0.5 if re.search(r"[-_][a-z]+[-_]", p) else 0)
+    if re.search(r"\.(php|aspx?)$", p) and len(segs) <= 1:
+        s -= 0.5
+    if re.search(r"/(privacy|terms|legal|cookie|imprint|impressum|contact|careers?|jobs|sitemap)", p, re.I):
+        s -= 1.5
+    return max(0.1, s)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--index", required=True, help="glob of round-1 crawl index files")
+    ap.add_argument("--extra", default="", help="curated URL list to add")
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--second", type=float, default=0.3)
+    args = ap.parse_args()
+    rnd = random.Random(7)
+    out, seen = [], set()
+    n = {"light": 0, "native": 0}
+    for f in sorted(glob.glob(args.index)):
+        for line in open(f):
+            try:
+                r = json.loads(line)
+            except Exception:
+                continue
+            kind = r.get("kind")
+            if kind not in ("light", "native_media", "native_class") or not r.get("links"):
+                continue
+            home = (r.get("final_url") or r["url"]).rstrip("/")
+            links = [l for l in r["links"] if l.rstrip("/") != home and l not in seen]
+            if not links:
+                continue
+            k = 2 if (kind == "light" and rnd.random() < args.second and len(links) > 1) else 1
+            for _ in range(k):
+                w = [score(l) for l in links]
+                pick = rnd.choices(links, weights=w)[0]
+                links.remove(pick)
+                seen.add(pick)
+                out.append(pick)
+                n["light" if kind == "light" else "native"] += 1
+                if not links:
+                    break
+    if args.extra:
+        for u in open(args.extra):
+            u = u.strip()
+            if u and not u.startswith("#") and u not in seen:
+                seen.add(u)
+                out.append(u)
+    rnd.shuffle(out)
+    open(args.out, "w").write("\n".join(out) + "\n")
+    print(len(out), n)
+
+
+if __name__ == "__main__":
+    main()

training/model.py

@@ -5,15 +5,24 @@ Attention logits get a learned per-head bias from the page's structure graph (te
 surface-on-surface, border-on-surface, same-element, image-on-surface, same-colour), so the model
 reasons about hierarchy and contrast the way a designer would. Outputs a dark colour per token
 (OKLab, parameterised relative to the input so hue is kept unless the model chooses otherwise),
-an alpha, and a treatment class per image.
+an alpha, and a compact page context.
+
+Element head (v2): decides how every image, CSS background, sprite or pseudo-element icon on the page is shown
+(keep / invert / adapt / remove) and what it is (photo, logo, icon, formula...). It reuses the trunk's image
+encoder and reads the page context plus the light surface under the element and the dark surface the palette
+head designed for it, so it can run for any element at any time (late-loaded icons, sprites) in a tiny second
+graph without re-running the transformer.
 """
 import math
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-FT, FI, FP = 49, 39, 8
-N_ROLES, N_TREAT = 7, 5
+FT, FI, FP = 49, 41, 8
+N_ROLES = 7
+TREAT = ["keep", "invert", "adapt", "remove"]
+IMG_ROLES = ["photo", "illustration", "diagram", "logo", "icon", "text", "backdrop", "avatar", "ui"]
+N_TREAT, N_IROLE, N_CTX = len(TREAT), len(IMG_ROLES), 32
 K_REL = 11
 
 
@@ -89,7 +98,19 @@ class NadirNet(nn.Module):
         self.blocks = nn.ModuleList([Block(d, heads, ff, drop) for _ in range(layers)])
         self.ln = nn.LayerNorm(d)
         self.color_head = MLP(d + 3, d, 6)
-        self.img_head = MLP(d, d // 2, N_TREAT)
+        self.ctx_head = nn.Linear(d, N_CTX)
+        self.elem_head = MLP(d + N_CTX + 72, d, N_TREAT + N_IROLE, drop=drop)
+
+    def embed_images(self, img_f, img_thumb):
+        return self.img_in(torch.cat([img_f, self.cnn(img_thumb)], -1))
+
+    def elements(self, e_f, e_thumb, under_light, under_dark, ctx):
+        """e_f (B,K,FI), e_thumb (B,K,4,32,32), under_light/under_dark (B,K,3) OKLab, ctx (B,N_CTX)
+        -> treatment logits (B,K,N_TREAT), role logits (B,K,N_IROLE)."""
+        K = e_f.shape[1]
+        h = torch.cat([self.embed_images(e_f, e_thumb), ctx[:, None].expand(-1, K, -1), fourier_lab(under_light), fourier_lab(under_dark)], -1)
+        z = self.elem_head(h)
+        return z[..., :N_TREAT], z[..., N_TREAT:]
 
     def forward(self, tok_f, tok_role, tok_lab, img_f, img_thumb, page_f, src, rel, mask):
         """tok_f (B,N,FT) tok_role (B,N) long, tok_lab (B,N,3), img_f (B,M,FI), img_thumb (B,M,4,32,32) float 0-1,
@@ -98,7 +119,7 @@ class NadirNet(nn.Module):
         M = img_f.shape[1]
         tok = self.tok_in(torch.cat([tok_f, F.one_hot(tok_role, N_ROLES).to(tok_f.dtype), fourier_lab(tok_lab)], -1)) + self.typ.weight[1]
         if M > 0:
-            img = self.img_in(torch.cat([img_f, self.cnn(img_thumb)], -1)) + self.typ.weight[2]
+            img = self.embed_images(img_f, img_thumb) + self.typ.weight[2]
         else:
             img = tok.new_zeros(B, 0, tok.shape[-1])
         cls = self.page_in(page_f)[:, None] + self.typ.weight[0]
@@ -117,8 +138,18 @@ class NadirNet(nn.Module):
         b = r * a_in + s * b_in + 0.12 * torch.tanh(z[..., 4])
         lab = torch.stack([L, a, b], -1)
         alpha_logit = z[..., 5]
-        img_logits = self.img_head(x[:, 1 + N:]) if M > 0 else x.new_zeros(B, 0, N_TREAT)
-        return lab, alpha_logit, img_logits
+        ctx = self.ctx_head(x[:, 0])
+        return lab, alpha_logit, ctx
+
+
+class ElemNet(nn.Module):
+    """The element head as its own graph (shares NadirNet's weights): classify any set of page images."""
+    def __init__(self, net):
+        super().__init__()
+        self.net = net
+
+    def forward(self, e_f, e_thumb, under_light, under_dark, ctx):
+        return self.net.elements(e_f, e_thumb, under_light, under_dark, ctx)
 
 
 # ---------- colour maths (differentiable) ----------

training/predict.py

@@ -9,10 +9,11 @@ from model import NadirNet
 from train import collate, prep
 from mkjobs import job_from
 
-TREAT = ["keep", "dim", "invert", "light", "plate"]
+TREAT = ["keep", "invert", "adapt", "remove"]
+RENDER_T = {"keep": "keep", "invert": "invert", "adapt": "invert", "remove": "clear"}  # render.py has no adapt path
 
 
-from safety import oklab, lab_to_rgb, contrast, safe_treatment, enforce_contrast, is_media_keep, hue_guard, keep_invisible
+from safety import oklab, lab_to_rgb, contrast, safe_treatment, enforce_contrast, is_media_keep, hue_guard, keep_invisible, rescue_surfaces, glass_panels, opaque_surfaces
 
 
 def decode(L, lab, alpha_logit, img_logits, rels):
@@ -29,11 +30,14 @@ def decode(L, lab, alpha_logit, img_logits, rels):
             out[i] = {"rgb": rgb, "a": float(t["rgba"][3]), "L": oklab(*rgb)[0]}
         else:
             hue_guard(t, out[i])
+    glass_panels(toks, out, keep)
     for i, t in enumerate(toks):
         if t["f"][24] > 0.5 and t["role"] == 0 and out[i]["L"] > 0.35:
             o = oklab(*out[i]["rgb"]); out[i]["rgb"] = lab_to_rgb(0.2, o[1] * 0.3, o[2] * 0.3); out[i]["L"] = 0.2
+    opaque_surfaces(toks, out, keep)
     keep_invisible(toks, out, rels, keep)
     enforce_contrast(toks, out, rels, keep)
+    rescue_surfaces(toks, out, rels, keep)
     def hx(o):
         h = "#%02x%02x%02x" % tuple(o["rgb"])
         return h + ("%02x" % round(o["a"] * 255) if o["a"] < 0.999 else "")
@@ -41,11 +45,24 @@ def decode(L, lab, alpha_logit, img_logits, rels):
     treat = {}
     for j, im in enumerate(L["images"][:32]):
         t = TREAT[int(np.argmax(img_logits[j]))] if j < len(img_logits) else "keep"
-        u = im.get("underTok", -1)
-        treat[str(j)] = safe_treatment(t, im, out[u]["L"] if u is not None and 0 <= u < len(out) else None)
+        treat[str(j)] = RENDER_T[t]
     return colors, sorted(keep), treat
 
 
+def element_treatments(model, s, lab, ctx, dev):
+    """Element head over a page's first 32 images (render jobs only know those)."""
+    m = min(32, len(s["IF"]))
+    if not m:
+        return []
+    f = torch.from_numpy(s["IF"][:m])[None].to(dev)
+    th = torch.from_numpy(s["thumbs"][:m].transpose(0, 3, 1, 2) / 255.0).float()[None].to(dev)
+    ul = torch.from_numpy(s["IF"][:m, 23:26] * np.array([1, 0.25, 0.25], np.float32))[None].to(dev)
+    u = np.where((s["img_under"][:m] >= 0) & (s["img_under"][:m] < lab.shape[1]), s["img_under"][:m], s["root"])
+    ud = lab[:, torch.from_numpy(u).to(dev)]
+    t, _ = model.elements(f, th, ul, ud, ctx)
+    return t[0].float().cpu().numpy().tolist()
+
+
 def main():
     ap = argparse.ArgumentParser()
     ap.add_argument("--ckpt", required=True)
@@ -70,10 +87,11 @@ def main():
             L = json.load(gzip.open(Path(args.crawl) / "pages" / s["id"] / "light.json.gz"))
             b = collate([s], dev)
             with torch.no_grad():
-                lab, al, il = model(b["tok_f"], b["role"], b["lab"], b["img_f"], b["thumbs"], b["page"], torch.zeros_like(b["src"]), b["rel"], b["mask"])
+                lab, al, ctx = model(b["tok_f"], b["role"], b["lab"], b["img_f"], b["thumbs"], b["page"], torch.zeros_like(b["src"]), b["rel"], b["mask"])
+                il = element_treatments(model, s, lab.float(), ctx, dev)
             n = len(s["role"])
             colors, keep, treat = decode(L, lab[0, :n].float().cpu().numpy().tolist(), al[0, :n].float().cpu().numpy().tolist(),
-                                         il[0].float().cpu().numpy().tolist(), s["rel"])
+                                         il, s["rel"])
             colors = colors + [None] * (len(L["tokens"]) - len(colors))
             j = job_from(L, s["id"], args.tag, colors, treat)
             j["keep"] = sorted(set(j["keep"]) | set(keep))

training/quantize.py

@@ -13,7 +13,7 @@ import torch.nn as nn
 
 from model import NadirNet
 from train import prep
-from export import feeds_for, fidelity, bench, NAMES_IN
+from export import feeds_for, elem_feeds, fidelity, bench, NAMES_IN
 
 GS = 32
 
@@ -245,25 +245,29 @@ def build_onnx(fp32_path, mapping, packed, out_path, accuracy_level=0):
     onnx.save(m, out_path)
 
 
-def quality(sess, samples):
-    """dE vs the training targets (Haiku for src 0, native dark for src 1) and image-treatment accuracy."""
+def quality(sess, samples, esess=None, ref=None):
+    """dE vs the training targets (Haiku for src 0, native dark for src 1); with the element graph, how often the
+    element decisions computed from this page graph's outputs match those from the fp32 page graph."""
     agg = {0: [0.0, 0.0], 1: [0.0, 0.0]}
     ok = n = 0
     for s in samples:
         f = feeds_for(s)
-        lab, _, imgl = sess.run(None, f)
+        lab, _, ctx = sess.run(None, f)
         k = len(s["role"])
         m = s["tmask"][:k]
         if m.any():
             d = np.linalg.norm(lab[0, :k] - s["lab_t"][:k], axis=-1)
             w = s["w"][:k] * m
             agg[int(s["src"])][0] += float((d * w).sum() / max(w.sum(), 1e-9)); agg[int(s["src"])][1] += 1
-        mi = len(s["IF"])
-        if mi:
-            t = np.asarray(s["imgt"][:mi])
-            v = t >= 0
-            ok += int((imgl[0, :mi].argmax(-1)[v] == t[v]).sum()); n += int(v.sum())
-    return {"haiku_dE": agg[0][0] / max(1, agg[0][1]), "native_dE": agg[1][0] / max(1, agg[1][1]), "img_acc": ok / max(1, n)}
+        if esess is not None and ref is not None and len(s["E_ut"]):
+            rl, _, rc = ref.run(None, f)
+            a = esess.run(None, elem_feeds(s, lab, ctx)[0])[0][0].argmax(-1)
+            b = esess.run(None, elem_feeds(s, rl, rc)[0])[0][0].argmax(-1)
+            ok += int((a == b).sum()); n += len(a)
+    out = {"haiku_dE": agg[0][0] / max(1, agg[0][1]), "native_dE": agg[1][0] / max(1, agg[1][1])}
+    if n:
+        out["elem_decisions_same_as_fp32"] = ok / n
+    return out
 
 
 def main():
@@ -275,6 +279,7 @@ def main():
     ap.add_argument("--calib", type=int, default=256)
     ap.add_argument("--eval", type=int, default=400)
     ap.add_argument("--threads", type=int, default=8)
+    ap.add_argument("--elem", default="", help="element graph (fp32) to measure decision agreement")
     args = ap.parse_args()
     torch.set_num_threads(args.threads)
     os.makedirs(args.out, exist_ok=True)
@@ -291,7 +296,7 @@ def main():
     mapping = onnx_matmuls(model, onnx.load(args.fp32))
     quantizable = [n for n in order if n in mapping and not n.endswith(".rel")]
     print(f"{len(quantizable)}/{len(order)} linears map to MatMul; fp32 kept: {[n for n in order if n not in quantizable]}", flush=True)
-    edges = [n for n in quantizable if n.startswith(("tok_in.", "img_in.", "color_head.", "img_head."))]
+    edges = [n for n in quantizable if n.startswith(("tok_in.", "img_in.", "color_head.", "ctx_head."))]
     plans = {"int8": {n: 8 for n in quantizable},
              "int4": {n: (8 if n in edges else 4) for n in quantizable}}
     so = ort.SessionOptions(); so.intra_op_num_threads = 4
@@ -299,7 +304,9 @@ def main():
     report = {"format": {"8-bit layers": f"uint8 blocks of {GS} with fp32 scale + uint8 zero point, dequantized in-graph (Cast/Sub/Mul, folded at load)", "4-bit layers": f"uint4 blocks of {GS}, MatMulNBits"},
               "group_size": GS, "calib_pages": len(calib), "eval_pages": len(val), "fp32_layers": [n for n in order if n not in quantizable],
               "int4_8bit_edge_layers": edges, "variants": {}}
-    r = quality(ref, val); r["size_mb"] = round(os.path.getsize(args.fp32) / 1e6, 2); r["cpu_ms_median"] = round(bench(ref, val), 2)
+    esess = ort.InferenceSession(args.elem, so, providers=["CPUExecutionProvider"]) if args.elem else None
+    bfeeds = [feeds_for(s) for s in val[:20]]
+    r = quality(ref, val); r["size_mb"] = round(os.path.getsize(args.fp32) / 1e6, 2); r["cpu_ms_median"] = round(bench(ref, bfeeds), 2)
     report["variants"]["fp32"] = r
     print("fp32", r, flush=True)
     for vname, bits_for in plans.items():
@@ -316,10 +323,10 @@ def main():
                 tq = qm(*[torch.from_numpy(f[k]) for k in NAMES_IN])[0].numpy()
                 par = max(par, float(np.abs(tq - sess.run(None, f)[0]).max()))
         r = fidelity(sess, ref, val)
-        r.update(quality(sess, val))
+        r.update(quality(sess, val, esess, ref))
         r["pack_parity_max_abs"] = par
         r["size_mb"] = round(os.path.getsize(path) / 1e6, 2)
-        r["cpu_ms_median"] = round(bench(sess, val), 2)
+        r["cpu_ms_median"] = round(bench(sess, bfeeds), 2)
         report["variants"][vname] = r
         print(vname, r, flush=True)
         torch.save({"cfg": ck["cfg"], "model": qm.state_dict()}, os.path.join(args.out, f"fakequant-{vname}.pt"))

training/safety.py

@@ -44,7 +44,9 @@ def adaptable(im, s, surf_L=None):
         return False
     if not s:
         return im["w"] * im["h"] <= 250000
-    return s["transp"] < 0.05 and s.get("ringA", 0) >= 0.95 and s.get("ringL", 0) >= 0.85 and s.get("white", 0) >= 0.3 and s["colors"] < 0.2
+    # on its own white rectangle, or on transparency with dark ink that would vanish (the pixel check confirms)
+    return (s["transp"] < 0.05 and s.get("ringA", 0) >= 0.95 and s.get("ringL", 0) >= 0.85 and s.get("white", 0) >= 0.3 and s["colors"] < 0.2) or \
+        (s["transp"] > 0.15 and s["meanL"] < 0.6)
 
 
 def is_light_backdrop(im, s):
@@ -69,7 +71,9 @@ def _base_treatment(t, im, surf_L=None):
     if is_light_backdrop(im, s):
         return "clear"
     # a mostly dark graphic on transparency would vanish on a dark surface: hue-preserving inversion
-    if (not t or t == "keep") and s and s["transp"] > 0.15 and s["meanL"] < 0.35 and im["w"] * im["h"] < 200000 and (surf_L is None or surf_L < 0.4):
+    # (icon-sized ones already when mid-grey: grey UI icons are as unreadable on dark as black ones)
+    if (not t or t == "keep") and s and s["transp"] > 0.15 and im["w"] * im["h"] < 200000 and (surf_L is None or surf_L < 0.4) and \
+            (s["meanL"] < 0.35 or (s["meanL"] < 0.5 and im["w"] * im["h"] <= 4096)):
         return "invert"
     if not t or t == "keep":
         return "keep"
@@ -108,6 +112,34 @@ def hue_guard(t, o):
         o["rgb"] = lab_to_rgb(q[0], q[1] * 0.06 / cout, q[2] * 0.06 / cout)
 
 
+def glass_panels(toks, out, keep):
+    """A see-through light panel (frosted card, glass button: white at 40-99% opacity) stays see-through but turns
+    dark; left mid-grey it is a surface no text colour reads on (enforce_contrast only sees opaque surfaces)."""
+    for i, t in enumerate(toks):
+        a = t["rgba"][3]
+        if i in keep or t["role"] != 0 or a >= 0.999 or a < 0.4:
+            continue
+        if oklab(*[int(x) for x in t["rgba"][:3]])[0] < 0.85 or out[i]["L"] <= 0.45:
+            continue
+        o = oklab(*out[i]["rgb"])
+        out[i]["rgb"] = lab_to_rgb(0.24, o[1] * 0.5, o[2] * 0.5)
+        out[i]["L"] = 0.24
+
+
+def opaque_surfaces(toks, out, keep):
+    """An opaque surface stays opaque: see-through, it would show whatever the design stacks beneath it (a hidden
+    sticky header, the next slide). The predicted colour is composited over the dark canvas instead."""
+    ci = next((i for i, t in enumerate(toks) if t["f"][24] > 0.5 and t["role"] == 0), -1)
+    under = out[ci]["rgb"] if ci >= 0 else [18, 18, 20]
+    for i, t in enumerate(toks):
+        if i in keep or t["role"] != 0 or t["rgba"][3] < 0.999 or out[i]["a"] >= 0.999:
+            continue
+        a = out[i]["a"]
+        out[i]["rgb"] = [int(round(c * a + u * (1 - a))) for c, u in zip(out[i]["rgb"], under)]
+        out[i]["a"] = 1.0
+        out[i]["L"] = oklab(*out[i]["rgb"])[0]
+
+
 def keep_invisible(toks, out, rels, keep):
     """Things that were invisible against their surface in the light design stay invisible: a gradient stop
     with exactly a surface's colour (text fades, edge fades) takes that surface's dark colour, and a border that
@@ -163,22 +195,20 @@ def is_media_keep(t):
 def enforce_contrast(toks, out, rels, keep):
     """Keep every foreground colour readable on all the surfaces it sits on, never asking for more contrast
     than the light design had. One lightness per token (tokens are shared), chosen on a grid to minimise the
-    weighted log-contrast deficit, with a small pull towards the model's own lightness."""
+    area-weighted log-contrast deficit plus the distance from the model's own lightness: the surfaces a colour
+    mostly sits on decide, and a few bright accent surfaces cannot drag it to a mid-grey that reads nowhere
+    (those surfaces are handled by rescue_surfaces)."""
     n = len(toks)
     surf = {}
     for typ, a, b, w in rels:
         a, b = int(a), int(b)
         if int(typ) == 0 and a < n and b < n and toks[b]["role"] == 0 and out[b]["a"] >= 0.999:
-            surf.setdefault(a, []).append((b, float(w) + 0.1))  # log-area weights: small buttons still count
+            surf.setdefault(a, []).append((b, math.expm1(float(w)) + 1.0))  # rels store log1p(weight)
     for i, lst in surf.items():
         t = toks[i]
         if i in keep or t["role"] not in FG_ROLES or out[i]["a"] < 0.999:
             continue
-        if t["role"] == 1:
-            fs, bold = t["f"][21] * 32, t["f"][23] > 0.5
-            need0 = 3.0 if (fs >= 24 or (fs >= 18.5 and bold)) else 4.5
-        else:
-            need0 = 3.0
+        need0 = _need(t)
         reqs = []
         for j, w in lst:
             need = min(need0, contrast(t["rgba"][:3], toks[j]["rgba"][:3]))
@@ -196,7 +226,49 @@ def enforce_contrast(toks, out, rels, keep):
         for k in range(51):
             Lx = k / 50
             rgb = lab_to_rgb(Lx, o[1], o[2])
-            c = deficit(rgb) / W + 0.05 * abs(Lx - o[0])
+            c = deficit(rgb) / W + abs(Lx - o[0])
             if best is None or c < best[0]:
                 best = (c, rgb, Lx)
         out[i]["rgb"], out[i]["L"] = best[1], best[2]
+
+
+def _need(t):
+    """WCAG target for a foreground token: 3:1 for large/bold text and graphics, 4.5:1 for body text."""
+    if t["role"] == 1:
+        fs, bold = t["f"][21] * 32, t["f"][23] > 0.5
+        return 3.0 if (fs >= 24 or (fs >= 18.5 and bold)) else 4.5
+    return 3.0
+
+
+def rescue_surfaces(toks, out, rels, keep):
+    """A light surface whose text no longer reads (its text colour is shared with the dark page and follows it)
+    is darkened within its own hue until that text reads again: a yellow note becomes dark gold, not a glare."""
+    n = len(toks)
+    on = {}
+    for typ, a, b, w in rels:
+        a, b = int(a), int(b)
+        if int(typ) == 0 and a < n and b < n and toks[b]["role"] == 0 and toks[a]["role"] in FG_ROLES:
+            on.setdefault(b, []).append((a, math.expm1(float(w)) + 1.0))
+    for j, lst in on.items():
+        s = toks[j]
+        if j in keep or s["f"][24] > 0.5 or out[j]["a"] < 0.999 or out[j]["L"] < 0.6:
+            continue
+        reqs = []
+        for i, w in lst:
+            if i in keep or out[i]["a"] < 0.999:
+                continue
+            need = min(_need(toks[i]), contrast(toks[i]["rgba"][:3], s["rgba"][:3]))
+            if need >= 1.5:
+                reqs.append((out[i]["rgb"], need, w))
+        if not reqs:
+            continue
+        W = sum(r[2] for r in reqs)
+        fail = lambda rgb: sum(w for fg, need, w in reqs if contrast(fg, rgb) < need) / W
+        if fail(out[j]["rgb"]) <= 0.5:
+            continue
+        o = oklab(*out[j]["rgb"])
+        for k in range(int(o[0] * 50), 9, -1):
+            rgb = lab_to_rgb(k / 50, o[1], o[2])
+            if fail(rgb) <= 0.2:
+                out[j]["rgb"], out[j]["L"] = rgb, k / 50
+                break

training/train.py

@@ -1,12 +1,12 @@
 #!/usr/bin/env python3
-"""Train NadirNet on Haiku-designed and native dark-mode targets."""
+"""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, contrast_ratio
+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))
@@ -35,6 +35,10 @@ def prep(s):
     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
 
 
@@ -57,6 +61,8 @@ def augment(s, rng):
             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:
@@ -70,6 +76,12 @@ def augment(s, rng):
         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
 
 
@@ -81,7 +93,10 @@ def collate(batch, device):
     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, 39), np.float32); thumbs = np.zeros((B, M, 4, 32, 32), np.float32); imgt = -np.ones((B, M), np.int64)
+    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)
@@ -93,6 +108,13 @@ def collate(batch, device):
             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):
@@ -119,12 +141,12 @@ def collate(batch, device):
         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 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)
-    t = lambda x: torch.from_numpy(x).to(device, non_blocking=True)
-    return dict(tok_f=t(tok_f), role=t(role), lab=t(lab), lab_t=t(lab_t), tmask=t(tmask), w=t(w), alpha_in=t(alpha_in), alpha_t=t(alpha_t),
-                img_f=t(img_f), thumbs=t(thumbs), imgt=t(imgt), page=t(page), src=t(src), rel=t(rel), mask=t(mask), tb_pair=t(tb_pair))
+        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):
@@ -152,9 +174,16 @@ def huber(x, d):
     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, img_logits = 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(); img_logits = img_logits.float()
+    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)
@@ -178,17 +207,27 @@ def losses(model, b, img_w=None):
         lcon = (F.relu(c_t - c_pred) * pm.float()).sum() / pm.float().sum().clamp_min(1)
     else:
         lcon = lab.new_zeros(())
-    # images
-    im = b["imgt"] >= 0
-    if im.any():
-        li = F.cross_entropy(img_logits[im], b["imgt"][im], weight=img_w, label_smoothing=0.05)
-        acc = (img_logits[im].argmax(-1) == b["imgt"][im]).float().mean()
+    # 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:
-        li = lab.new_zeros(()); acc = lab.new_zeros(())
+        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, li, dict(dE=de.item(), dL=dl.item(), img_acc=acc.item(), n_img=int(im.sum()))
+    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):
@@ -199,16 +238,22 @@ def evaluate(model, val, device, bs, img_w):
             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, li, st = losses(model, b, img_w)
+                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_img"]:
-                    agg["img_acc"] += st["img_acc"] * st["n_img"]; cnt["img"] += st["n_img"]
+                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 == "img_acc":
-            out[k] = v / max(1, cnt["img"])
+        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()
@@ -230,25 +275,42 @@ def main():
     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"
+    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)
-    # image class weights (inverse sqrt frequency)
-    cnt = np.ones(5)
+    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:
-        for v in s["imgt"]:
-            if v >= 0:
-                cnt[v] += 1
-    img_w = torch.tensor((cnt.sum() / cnt) ** 0.5, dtype=torch.float32, device=device)
-    img_w = img_w / img_w.mean()
-    print("image label counts", cnt.astype(int).tolist(), flush=True)
+        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
@@ -266,8 +328,8 @@ def main():
         for bnp in loader:
             b = to_dev(bnp, device)
             with torch.autocast("cuda", dtype=torch.bfloat16):
-                lc, la, lcon, li, st = losses(model, b, img_w)
-            loss = lc + 0.5 * la + 0.02 * lcon + 0.05 * li
+                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
@@ -275,16 +337,42 @@ def main():
             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["img"] += li.item(); nb += 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.02 * (1 - ev.get("img_acc", 0))
+        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)