ProCreations commited on
Commit
bde89e0
·
verified ·
1 Parent(s): dd83935

Nadir 1.0: 15.9M on-device dark-mode model, int8/int4 quants, runtime, extension

Browse files
README.md ADDED
@@ -0,0 +1,218 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: apache-2.0
3
+ library_name: onnx
4
+ pipeline_tag: other
5
+ tags:
6
+ - dark-mode
7
+ - browser-extension
8
+ - onnx
9
+ - onnxruntime-web
10
+ - accessibility
11
+ - design
12
+ - distillation
13
+ ---
14
+
15
+ # Nadir
16
+
17
+ **Nadir** is a small on-device model that gives light websites a tasteful, readable dark mode, with the
18
+ care of a real dark theme rather than a filter. It runs entirely in your browser (onnxruntime-web, WASM),
19
+ in about 35 ms of model time per page, and ships with a ready-to-load Chromium extension.
20
+
21
+ Nadir follows three rules:
22
+
23
+ 1. **Already dark? Leave it alone.** Pages that are already dark are detected and never touched.
24
+ 2. **Has its own dark mode? Use it.** Nadir switches on the site's own dark theme (`prefers-color-scheme`
25
+ media rules, or `dark`/`theme-dark` class and attribute toggles it detects in the page's CSS).
26
+ 3. **No dark mode? The model designs one.** Nadir reads the page's colors and images and predicts a dark
27
+ palette for that specific site: brand colors stay recognizable, surfaces keep their hierarchy, text stays
28
+ readable, and logos stay visible: a logo on its own white rectangle is adapted to sit cleanly on the dark
29
+ surface. Photos, videos and ads are never color-shifted.
30
+
31
+ In a blind pairwise test on 146 sites it never trained on, judged by Claude Haiku 5.5, the preference
32
+ went to Nadir's dark mode over Dark Reader's on **93 sites vs 50** (3 ties): 65% of the decided pages.
33
+
34
+ ## How it works
35
+
36
+ Nadir doesn't paint pixels. It redesigns the page's **color system**:
37
+
38
+ 1. **Extraction (`runtime/extract.js`).** One synchronous pass over the DOM collects every distinct color
39
+ *in its role*: background, text, border, shadow, gradient, SVG fill, SVG stroke. Each token gets 49
40
+ features: lightness, chroma, hue, area, text size and weight, where it appears (header, nav, link, form,
41
+ code, footer), what it sits on and the original contrast. It also records four relations between tokens
42
+ (text-on-surface, surface-on-surface, border-on-surface, same-element) and every visible image or
43
+ inline-SVG logo with 39 features plus a 32×32 thumbnail. Ad slots are excluded.
44
+ 2. **Model (`NadirNet`).** A relation-biased transformer over `[page] + tokens + images`. Each attention head
45
+ gets a learned bias from the 11 relation channels, so the model knows which text sits on which surface.
46
+ It predicts for every token an OKLab color as a parameterized edit of the input (lightness, chroma scale,
47
+ hue rotation, tint, alpha). This keeps brand hues stable and makes impossible colors rare. For every
48
+ image it predicts a treatment: `keep`, `dim`, `invert`, `light` (monochrome dark glyph → near-white) or
49
+ `plate` (dark multicolor logo that needs help; shown with a hue-preserving inversion at runtime). A small
50
+ CNN reads the thumbnails.
51
+ 3. **Design rules (`runtime/nadir-infer.js`, `runtime/nadir-map.js`).** Deterministic guards on top of the
52
+ model:
53
+ - The page canvas is always dark.
54
+ - Every text and icon color keeps WCAG contrast on *every* surface it actually sits on. It never needs
55
+ more contrast than the original design had.
56
+ - Brand colors keep their hue (±45°).
57
+ - Things that were invisible in the light design stay invisible: text-fade gradients and borders that
58
+ matched their surface follow the surface.
59
+ - Colors drawn over video and photos are left alone when they already read as dark mode (player
60
+ controls, scrims, captions).
61
+ - Photos and video are never recolored.
62
+ - Light, flat background images behind content (paper textures, pale washes, SVG gradients) are cleared so the
63
+ designed dark surface shows through. Images with real detail are left alone.
64
+ - Inline SVG logos stay readable: their shapes are drawn as ink, and SVG mask and clip-path shapes (knockout
65
+ text, cut-outs) are never recolored.
66
+ - **Logos on their own white rectangle are adapted.** The white background is removed so Nadir's surface
67
+ shows through, black or gray ink becomes light, and brand colors keep their hue (lifted only as far as
68
+ needed to read on dark). This applies only to images the page itself marks as its logo.
69
+ - Dark logos on transparency get a hue-preserving inversion, and already-light ones are kept.
70
+ 4. **Engine.** The predicted palette is applied through a patched
71
+ [Dark Reader](https://github.com/darkreader/darkreader) 4.9.133 dynamic engine (MIT), which handles the
72
+ hard plumbing: stylesheets, inline styles, shadow DOM and dynamically inserted content. It asks Nadir for
73
+ every color it rewrites. Colors the model didn't see (new content) are mapped by nearest-neighbor
74
+ lookup in the predicted palette.
75
+
76
+ ## Model
77
+
78
+ | | |
79
+ |---|---|
80
+ | Parameters | 15.9M (15,895,995) |
81
+ | Width / depth / heads | 384 / 8 layers / 6 heads |
82
+ | Inputs | ≤256 color tokens, ≤32 images, 11 relation channels |
83
+ | Outputs | OKLab color + alpha per token, 5-way image treatment |
84
+ | Runtime | onnxruntime-web 1.29 (WASM SIMD, multi-threaded) |
85
+
86
+ ### Files
87
+
88
+ | File | Precision | Size | Δ color vs fp32 (mean / max, OKLab) | Image agreement |
89
+ |---|---|---|---|---|
90
+ | `onnx/nadir-fp32.onnx` | fp32 | 63.7 MB | 0 / 0 | 100% |
91
+ | `onnx/nadir-int8.onnx` | 8-bit weights | 20.6 MB | 0.00015 / 0.025 | 99.99% |
92
+ | `onnx/nadir-int4.onnx` | 4-bit weights (8-bit input/output layers) | 13.3 MB | 0.00036 / 0.071 | 99.76% |
93
+
94
+ Fidelity is measured on 400 held-out pages. A difference of 0.02 in OKLab is about the smallest one you can see.
95
+
96
+ Speed in onnxruntime-web 1.29 (Chrome, Apple M4 Max, 4 WASM threads), per page:
97
+
98
+ | File | Load | Inference |
99
+ |---|---|---|
100
+ | fp32 | ~7 s | 34 ms |
101
+ | int8 | ~1.5 s | 33–35 ms |
102
+ | int4 | ~1.0 s | 38 ms |
103
+
104
+ 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
105
+ from cache.
106
+
107
+ - `onnx/nadir-int8.onnx` is the extension's default.
108
+ - The 8-bit and 4-bit models are **calibrated weight-only quants**:
109
+ - Layer-sequential GPTQ, with Hessians collected from 256 real crawled pages flowing through the
110
+ already-quantized earlier layers.
111
+ - Asymmetric codes in groups of 32, with each group's clipping range chosen by an MSE search.
112
+ - The 4-bit model keeps its input and output layers at 8 bits.
113
+ - Storage format:
114
+ - 4-bit layers use ONNX `MatMulNBits` (block 32).
115
+ - 8-bit layers store uint8 codes with a per-block scale and zero point, dequantized in the graph by standard
116
+ ops (`Cast`/`Sub`/`Mul`/`Reshape`). onnxruntime folds these into the weights when it loads the model, so
117
+ the 8-bit layers run as exact fp32 math with no special kernels. Both files are tested on onnxruntime 1.22
118
+ (CPU) and onnxruntime-web 1.29.
119
+ - Each quant is verified three ways: packing parity against a PyTorch fake-quant model (≤2e-6), fidelity
120
+ against fp32, and quality against the training targets.
121
+ - They run in onnxruntime-web with outputs matching onnxruntime CPU to ~1e-5.
122
+ - Activation-quantized int8 (ORT dynamic/static) was rejected: worst-case color error up to 0.99.
123
+
124
+ `model.safetensors` + `model.py` hold the PyTorch weights and architecture.
125
+
126
+ ## Training data
127
+
128
+ - **Pages.** A crawl of 47,569 sites from the [Tranco](https://tranco-list.eu/) top list at 1280×800.
129
+ Each page was captured in light mode, in emulated dark mode, and with dark-class toggles applied.
130
+ - **Teacher.** 19,563 light pages were redesigned by **Claude Haiku 5.5** (Anthropic Batch API,
131
+ high effort) acting as the site's own product designer. It was given the screenshot, a contact sheet of
132
+ the images, and the full color inventory with roles and locations, and returned a dark palette plus image
133
+ treatments in a strict JSON schema. The design brief is in `training/haiku_gen.py`.
134
+ - **Native pairs.** 4,142 pages with a real built-in dark theme give human-designed light→dark
135
+ pairs, aligned token by token.
136
+ - **Quality control.**
137
+ - Designs that missed ids, kept a light canvas, or left >15% of text below 3:1 contrast were dropped.
138
+ - The rest were passed through the same contrast rule the runtime uses (14,338 of 19,464 designs adjusted).
139
+ - The train/validation split is by site (hash of the hostname).
140
+ - **Augmentation.** Random token dropout (it also simulates ad blockers and partial loads) and whole-design
141
+ hue rotation.
142
+
143
+ ## Evaluation
144
+
145
+ The judge was a blind pairwise one: Claude Haiku 5.5 saw the original light page plus two dark versions in
146
+ random order, A and B. It judged readability first, then breakage (glaring light patches, lost logos,
147
+ distorted photos), then design quality. All pages are held-out sites the model never trained on.
148
+
149
+ | Nadir vs … (146 held-out sites) | Nadir preferred | Other preferred | Ties |
150
+ |---|---|---|---|
151
+ | Dark Reader 4.9 (dynamic mode) | **93** | 50 | 3 |
152
+ | Claude Haiku 5.5's own designs (the teacher)¹ | 76 | 61 | 9 |
153
+ | Nadir before its final design rules² | 45 | 29 | 72 |
154
+
155
+ ¹ The teacher's designs were rendered with the earlier runtime, so part of this margin comes from the runtime's newer
156
+ design rules rather than from the model itself. Before those rules, the model and its teacher were even (63 vs 67).
157
+ ² The rules for light background textures, SVG logos and SVG masks.
158
+
159
+ **Model size.** A 34.0M-parameter version (512 / 10 layers / 8 heads) trained on the same data was no better: 49 vs 53
160
+ with 45 ties against this one. It reached the same color error and runs about 2× slower in the browser, so the 15.9M
161
+ model is the release.
162
+
163
+ Held-out color error vs the teacher designs (OKLab ΔE): 0.040 (Haiku pages), 0.103
164
+ (native dark themes). Image-treatment accuracy: 79.4% over 5 classes.
165
+
166
+ ## Use it
167
+
168
+ ### Browser extension (Chrome, Edge, Brave, Arc, Aside and other Chromium browsers)
169
+
170
+ 1. Download `nadir-extension.zip` and unzip it.
171
+ 2. Open `chrome://extensions`, enable **Developer mode**, click **Load unpacked**, and choose the folder.
172
+ 3. Browse. The toolbar badge shows what Nadir did on each site:
173
+ - `AI`: model-designed dark mode.
174
+ - `ON`: the site's own dark mode.
175
+ - No badge: already dark.
176
+
177
+ The popup can turn Nadir off for a site.
178
+
179
+ The extension works alongside content blockers such as uBlock Origin Lite: removed ads simply aren't part
180
+ of the palette, and ads that remain are treated as images (never recolored).
181
+
182
+ ### In your own code
183
+
184
+ ```js
185
+ import * as ort from 'onnxruntime-web';
186
+ // runtime/extract.js, runtime/nadir-infer.js and runtime/nadir-map.js define the globals used below
187
+ const page = await __nadirExtract({describe: false, thumbs: false}); // in the page: colors, roles, relations, images
188
+ // pixel stats + thumbnails for the images (any context that can fetch them, e.g. an extension page)
189
+ const st = await __nadirImageStats(page.images, {dpr: devicePixelRatio});
190
+ page.images.forEach((im, i) => { Object.assign(im, st[i]); im.f = __nadirImageFeatures(im, page.vw, page.vh, page.docH); });
191
+ const built = NadirInfer.buildInputs(page);
192
+ const session = await ort.InferenceSession.create('nadir-int8.onnx', {executionProviders: ['wasm']});
193
+ const feeds = Object.fromEntries(Object.entries(built.feeds).map(([k, [data, dims]]) =>
194
+ [k, new ort.Tensor(k === 'tok_role' || k === 'src' ? 'int64' : k === 'mask' ? 'bool' : 'float32', data, dims)]));
195
+ const out = await session.run(feeds);
196
+ const {colors, keep, treatments} = NadirInfer.decode(page, built,
197
+ {lab: out.lab.data, alpha_logit: out.alpha_logit.data, img_logits: out.img_logits.data});
198
+ // colors[i] is the dark replacement for page.tokens[i] (null = keep); treatments[j] applies to page.images[j]
199
+ ```
200
+
201
+ ## Limitations
202
+
203
+ - Nadir recolors CSS. Text and UI baked into images, canvas drawings and cross-origin iframes keep their
204
+ original colors, though the iframes get their own Nadir pass in the extension.
205
+ - Colors are mapped per (color, role). If one exact color is used both as a light card and as a light
206
+ button, both get the same dark color.
207
+ - Background images that mix light areas with strong color (for example a colored hero shape with a white base) are
208
+ kept as they are, so their light parts stay light.
209
+ - The model reads at most 256 colors per page. On very colorful pages the rarest colors are mapped by nearest
210
+ neighbor in the predicted palette.
211
+ - The teacher is an LLM. Its taste is good but not perfect, and the model inherits its preferences.
212
+ - The judge is also an LLM, so the evaluation is automated, not a human study.
213
+
214
+ ## License
215
+
216
+ Apache-2.0 for the model, weights and Nadir code. The bundled page engine is a modified
217
+ [Dark Reader](https://github.com/darkreader/darkreader) (MIT License, © Dark Reader Ltd.); its license is
218
+ included in the extension (`lib/DARKREADER-LICENSE`).
config.json ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "d": 384,
3
+ "layers": 8,
4
+ "heads": 6,
5
+ "ff": 4,
6
+ "drop": 0.1,
7
+ "architecture": "NadirNet",
8
+ "token_features": 49,
9
+ "image_features": 39,
10
+ "relation_channels": 11,
11
+ "max_tokens": 256,
12
+ "max_images": 32,
13
+ "image_treatments": [
14
+ "keep",
15
+ "dim",
16
+ "invert",
17
+ "light",
18
+ "plate"
19
+ ],
20
+ "inputs": [
21
+ "tok_f",
22
+ "tok_role",
23
+ "tok_lab",
24
+ "img_f",
25
+ "img_thumb",
26
+ "page_f",
27
+ "src",
28
+ "rel",
29
+ "mask"
30
+ ],
31
+ "outputs": [
32
+ "lab",
33
+ "alpha_logit",
34
+ "img_logits"
35
+ ]
36
+ }
model.py ADDED
@@ -0,0 +1,141 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """NadirNet: a relation-biased set transformer over a page's colour palette and images.
2
+
3
+ Tokens: [PAGE] + one token per (colour, role) the page uses + one token per image.
4
+ Attention logits get a learned per-head bias from the page's structure graph (text-on-surface,
5
+ surface-on-surface, border-on-surface, same-element, image-on-surface, same-colour), so the model
6
+ reasons about hierarchy and contrast the way a designer would. Outputs a dark colour per token
7
+ (OKLab, parameterised relative to the input so hue is kept unless the model chooses otherwise),
8
+ an alpha, and a treatment class per image.
9
+ """
10
+ import math
11
+ import torch
12
+ import torch.nn as nn
13
+ import torch.nn.functional as F
14
+
15
+ FT, FI, FP = 49, 39, 8
16
+ N_ROLES, N_TREAT = 7, 5
17
+ K_REL = 11
18
+
19
+
20
+ def fourier_lab(lab, n=6):
21
+ x = torch.stack([lab[..., 0], lab[..., 1] * 4, lab[..., 2] * 4], -1) # (..., 3)
22
+ freqs = (2.0 ** torch.arange(n, device=lab.device, dtype=lab.dtype)) * math.pi
23
+ ang = x[..., None] * freqs # (...,3,n)
24
+ return torch.cat([torch.sin(ang), torch.cos(ang)], -1).flatten(-2) # (..., 6n)
25
+
26
+
27
+ class MLP(nn.Module):
28
+ def __init__(self, i, h, o, drop=0.0):
29
+ super().__init__()
30
+ self.net = nn.Sequential(nn.Linear(i, h), nn.GELU(), nn.Dropout(drop), nn.Linear(h, o))
31
+
32
+ def forward(self, x):
33
+ return self.net(x)
34
+
35
+
36
+ class ThumbCNN(nn.Module):
37
+ """32x32 RGBA thumbnail -> embedding. Alpha matters: transparent logos behave differently from photos."""
38
+ def __init__(self, o=128):
39
+ super().__init__()
40
+ c = [5, 32, 64, 96, 128]
41
+ layers = []
42
+ for a, b in zip(c[:-1], c[1:]):
43
+ layers += [nn.Conv2d(a, b, 3, padding=1), nn.GELU(), nn.Conv2d(b, b, 3, padding=1, stride=2), nn.GELU()]
44
+ self.net = nn.Sequential(*layers)
45
+ self.proj = nn.Linear(128 * 2 * 2, o)
46
+
47
+ def forward(self, x): # (B, M, 4, 32, 32) in [0,1]
48
+ B, M = x.shape[:2]
49
+ x = x.reshape(B * M, 4, 32, 32)
50
+ a = x[:, 3:4]
51
+ lum = (0.2126 * x[:, 0:1] + 0.7152 * x[:, 1:2] + 0.0722 * x[:, 2:3]) * a # premultiplied luminance
52
+ x = torch.cat([x[:, :3] * a, a, lum], 1)
53
+ x = self.net(x) # (BM,128,2,2)
54
+ return self.proj(x.flatten(1)).reshape(B, M, -1)
55
+
56
+
57
+ class Block(nn.Module):
58
+ def __init__(self, d, h, ff, drop):
59
+ super().__init__()
60
+ self.h = h
61
+ self.ln1 = nn.LayerNorm(d)
62
+ self.qkv = nn.Linear(d, 3 * d)
63
+ self.proj = nn.Linear(d, d)
64
+ self.rel = nn.Linear(K_REL, h, bias=False)
65
+ nn.init.normal_(self.rel.weight, std=0.5)
66
+ self.ln2 = nn.LayerNorm(d)
67
+ self.ff = nn.Sequential(nn.Linear(d, ff * d), nn.GELU(), nn.Dropout(drop), nn.Linear(ff * d, d))
68
+ self.drop = nn.Dropout(drop)
69
+
70
+ def forward(self, x, rel, pad_bias):
71
+ B, T, D = x.shape
72
+ q, k, v = self.qkv(self.ln1(x)).reshape(B, T, 3, self.h, D // self.h).permute(2, 0, 3, 1, 4)
73
+ bias = self.rel(rel).permute(0, 3, 1, 2) + pad_bias # (B,H,T,T)
74
+ a = F.scaled_dot_product_attention(q, k, v, attn_mask=bias.to(q.dtype))
75
+ x = x + self.drop(self.proj(a.transpose(1, 2).reshape(B, T, D)))
76
+ return x + self.drop(self.ff(self.ln2(x)))
77
+
78
+
79
+ class NadirNet(nn.Module):
80
+ def __init__(self, d=384, layers=8, heads=6, ff=4, drop=0.1):
81
+ super().__init__()
82
+ self.cfg = dict(d=d, layers=layers, heads=heads, ff=ff, drop=drop)
83
+ self.tok_in = MLP(FT + N_ROLES + 36, 2 * d, d)
84
+ self.cnn = ThumbCNN(128)
85
+ self.img_in = MLP(FI + 128, 2 * d, d)
86
+ self.page_in = MLP(FP, d, d)
87
+ self.src = nn.Embedding(2, d)
88
+ self.typ = nn.Embedding(3, d)
89
+ self.blocks = nn.ModuleList([Block(d, heads, ff, drop) for _ in range(layers)])
90
+ self.ln = nn.LayerNorm(d)
91
+ self.color_head = MLP(d + 3, d, 6)
92
+ self.img_head = MLP(d, d // 2, N_TREAT)
93
+
94
+ def forward(self, tok_f, tok_role, tok_lab, img_f, img_thumb, page_f, src, rel, mask):
95
+ """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,
96
+ page_f (B,FP), src (B,) long, rel (B,T,T,K) with T=1+N+M, mask (B,T) bool (True=real)."""
97
+ B, N = tok_role.shape
98
+ M = img_f.shape[1]
99
+ 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]
100
+ if M > 0:
101
+ img = self.img_in(torch.cat([img_f, self.cnn(img_thumb)], -1)) + self.typ.weight[2]
102
+ else:
103
+ img = tok.new_zeros(B, 0, tok.shape[-1])
104
+ cls = self.page_in(page_f)[:, None] + self.typ.weight[0]
105
+ x = torch.cat([cls, tok, img], 1) + self.src(src)[:, None]
106
+ pad_bias = torch.zeros(mask.shape, dtype=x.dtype, device=x.device).masked_fill(~mask, -1e4)[:, None, None, :]
107
+ for blk in self.blocks:
108
+ x = blk(x, rel, pad_bias)
109
+ x = self.ln(x)
110
+ ht = x[:, 1:1 + N]
111
+ z = self.color_head(torch.cat([ht, tok_lab], -1))
112
+ L = torch.sigmoid(z[..., 0])
113
+ s = 1.6 * torch.sigmoid(z[..., 1])
114
+ r = 0.6 * torch.tanh(z[..., 2])
115
+ a_in, b_in = tok_lab[..., 1], tok_lab[..., 2]
116
+ a = s * a_in - r * b_in + 0.12 * torch.tanh(z[..., 3])
117
+ b = r * a_in + s * b_in + 0.12 * torch.tanh(z[..., 4])
118
+ lab = torch.stack([L, a, b], -1)
119
+ alpha_logit = z[..., 5]
120
+ img_logits = self.img_head(x[:, 1 + N:]) if M > 0 else x.new_zeros(B, 0, N_TREAT)
121
+ return lab, alpha_logit, img_logits
122
+
123
+
124
+ # ---------- colour maths (differentiable) ----------
125
+ OKLAB_TO_LMS = torch.tensor([[1.0, 0.3963377774, 0.2158037573], [1.0, -0.1055613458, -0.0638541728], [1.0, -0.0894841775, -1.2914855480]])
126
+ LMS_TO_RGB = torch.tensor([[4.0767416621, -3.3077115913, 0.2309699292], [-1.2684380046, 2.6097574011, -0.3413193965], [-0.0041960863, -0.7034186147, 1.7076147010]])
127
+
128
+
129
+ def oklab_to_linrgb(lab):
130
+ lms = lab @ OKLAB_TO_LMS.to(lab).T
131
+ return (lms ** 3) @ LMS_TO_RGB.to(lab).T
132
+
133
+
134
+ def luminance(lab):
135
+ rgb = oklab_to_linrgb(lab).clamp(0, 1)
136
+ return rgb @ torch.tensor([0.2126, 0.7152, 0.0722]).to(lab)
137
+
138
+
139
+ def contrast_ratio(lab1, lab2):
140
+ l1, l2 = luminance(lab1), luminance(lab2)
141
+ return (torch.maximum(l1, l2) + 0.05) / (torch.minimum(l1, l2) + 0.05)
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e7acc89ba8a039a4889d89ec9f01afce49b2f765d1ac57a489eb4c604aef2ceb
3
+ size 63596788
nadir-extension.zip ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:4a5da0a2671306af61a162f81c8cfa276ab2dd70fb05f7e927a401d9e1f08757
3
+ size 23421769
onnx/nadir-fp32.onnx ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:83c54558313415596d4d3b3e753ff30e3a72cdc65e27a5c8eefcb132b9d5c934
3
+ size 63699911
onnx/nadir-int4.onnx ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:5e11864ad3d7f1edb7bc8b7006e5d7a1c12401bbc17578f590182c0ca764802a
3
+ size 13313059
onnx/nadir-int8.onnx ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:bbdb0a2b13bfca0c05f8d19d54d3b721fb4fcd459d40e8519a7a5f77af15d634
3
+ size 20626640
reports/eval_judge.json ADDED
@@ -0,0 +1,88 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "judge": "Claude Haiku 5.5 (Batch API, effort medium), blind pairwise: original light page + two dark versions in random order; criteria readability > breakage > design",
3
+ "pages": "held-out sites (hash split by hostname), rendered live at 1280x800",
4
+ "final_release_vs_others": {
5
+ "final_m3:dr": {
6
+ "n": 146,
7
+ "final_m3_wins": 93,
8
+ "dr_wins": 50,
9
+ "ties": 3,
10
+ "final_m3_winrate_excl_ties": 0.65,
11
+ "mean_score_final_m3": 6.85,
12
+ "mean_score_dr": 6.1
13
+ },
14
+ "final_m3:final_m": {
15
+ "n": 146,
16
+ "final_m3_wins": 45,
17
+ "final_m_wins": 29,
18
+ "ties": 72,
19
+ "final_m3_winrate_excl_ties": 0.608,
20
+ "mean_score_final_m3": 6.62,
21
+ "mean_score_final_m": 6.16
22
+ },
23
+ "final_m3:haiku": {
24
+ "n": 146,
25
+ "final_m3_wins": 76,
26
+ "haiku_wins": 61,
27
+ "ties": 9,
28
+ "final_m3_winrate_excl_ties": 0.555,
29
+ "mean_score_final_m3": 6.72,
30
+ "mean_score_haiku": 6.31
31
+ }
32
+ },
33
+ "before_final_design_rules": {
34
+ "final_m:dr": {
35
+ "n": 147,
36
+ "final_m_wins": 83,
37
+ "dr_wins": 57,
38
+ "ties": 7,
39
+ "final_m_winrate_excl_ties": 0.593,
40
+ "mean_score_final_m": 6.5,
41
+ "mean_score_dr": 6.28
42
+ },
43
+ "final_l:dr": {
44
+ "n": 147,
45
+ "final_l_wins": 81,
46
+ "dr_wins": 60,
47
+ "ties": 6,
48
+ "final_l_winrate_excl_ties": 0.574,
49
+ "mean_score_final_l": 6.61,
50
+ "mean_score_dr": 6.31
51
+ },
52
+ "final_m:final_l": {
53
+ "n": 147,
54
+ "final_m_wins": 49,
55
+ "final_l_wins": 53,
56
+ "ties": 45,
57
+ "final_m_winrate_excl_ties": 0.48,
58
+ "mean_score_final_m": 6.43,
59
+ "mean_score_final_l": 6.54
60
+ },
61
+ "final_m:haiku": {
62
+ "n": 146,
63
+ "final_m_wins": 63,
64
+ "haiku_wins": 67,
65
+ "ties": 16,
66
+ "final_m_winrate_excl_ties": 0.485,
67
+ "mean_score_final_m": 6.32,
68
+ "mean_score_haiku": 6.38
69
+ },
70
+ "final_l:haiku": {
71
+ "n": 146,
72
+ "final_l_wins": 64,
73
+ "haiku_wins": 70,
74
+ "ties": 12,
75
+ "final_l_winrate_excl_ties": 0.478,
76
+ "mean_score_final_l": 6.47,
77
+ "mean_score_haiku": 6.44
78
+ }
79
+ },
80
+ "notes": {
81
+ "final_m3": "released model (15.9M) with final runtime rules",
82
+ "final_m": "same model, runtime before the backdrop/SVG rules",
83
+ "final_l": "34.0M model, earlier runtime",
84
+ "dr": "Dark Reader 4.9.133 dynamic mode",
85
+ "haiku": "teacher designs from Claude Haiku 5.5, earlier runtime"
86
+ },
87
+ "cost_usd": 0.368
88
+ }
reports/export_report.json ADDED
@@ -0,0 +1,20 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "params": 15895995,
3
+ "parity_max_abs": 1.2516975402832031e-06,
4
+ "variants": {
5
+ "fp32": {
6
+ "dE_vs_fp32": 0.0,
7
+ "dE_max": 0.0,
8
+ "img_agree": 1.0,
9
+ "size_mb": 63.7,
10
+ "cpu_ms_median": 7.48
11
+ },
12
+ "fp16": {
13
+ "dE_vs_fp32": 0.0002421401860071152,
14
+ "dE_max": 0.23341728746891022,
15
+ "img_agree": 1.0,
16
+ "size_mb": 31.91,
17
+ "cpu_ms_median": 10.17
18
+ }
19
+ }
20
+ }
reports/ortweb_timing_m4max.json ADDED
@@ -0,0 +1,67 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "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",
3
+ "fp32": {
4
+ "t1": {
5
+ "isolated": false,
6
+ "load_ms": 7246.9,
7
+ "ms_median": 117.08,
8
+ "max_abs_lab_vs_ort_cpu": 6.16908073425293e-06,
9
+ "max_abs_img_vs_ort_cpu": 4.172325134277344e-06
10
+ },
11
+ "t4": {
12
+ "isolated": false,
13
+ "load_ms": 7363.2,
14
+ "ms_median": 33.93,
15
+ "max_abs_lab_vs_ort_cpu": 6.16908073425293e-06,
16
+ "max_abs_img_vs_ort_cpu": 4.172325134277344e-06
17
+ }
18
+ },
19
+ "int8": {
20
+ "t1": {
21
+ "isolated": false,
22
+ "load_ms": 1356.6,
23
+ "ms_median": 116.78,
24
+ "max_abs_lab_vs_ort_cpu": 1.9669532775878906e-06,
25
+ "max_abs_img_vs_ort_cpu": 2.9802322387695312e-06
26
+ },
27
+ "t4": {
28
+ "isolated": false,
29
+ "load_ms": 1628.3,
30
+ "ms_median": 34.78,
31
+ "max_abs_lab_vs_ort_cpu": 1.9669532775878906e-06,
32
+ "max_abs_img_vs_ort_cpu": 2.9802322387695312e-06
33
+ }
34
+ },
35
+ "int8_repeat": {
36
+ "t1": {
37
+ "isolated": false,
38
+ "load_ms": 1523.5,
39
+ "ms_median": 115.82,
40
+ "max_abs_lab_vs_ort_cpu": 1.9669532775878906e-06,
41
+ "max_abs_img_vs_ort_cpu": 2.9802322387695312e-06
42
+ },
43
+ "t4": {
44
+ "isolated": false,
45
+ "load_ms": 1519,
46
+ "ms_median": 32.68,
47
+ "max_abs_lab_vs_ort_cpu": 1.9669532775878906e-06,
48
+ "max_abs_img_vs_ort_cpu": 2.9802322387695312e-06
49
+ }
50
+ },
51
+ "int4": {
52
+ "t1": {
53
+ "isolated": false,
54
+ "load_ms": 1074.8,
55
+ "ms_median": 137.33,
56
+ "max_abs_lab_vs_ort_cpu": 1.3709068298339844e-06,
57
+ "max_abs_img_vs_ort_cpu": 7.659196853637695e-06
58
+ },
59
+ "t4": {
60
+ "isolated": false,
61
+ "load_ms": 1047,
62
+ "ms_median": 38.47,
63
+ "max_abs_lab_vs_ort_cpu": 1.3709068298339844e-06,
64
+ "max_abs_img_vs_ort_cpu": 7.659196853637695e-06
65
+ }
66
+ }
67
+ }
reports/quant_report.json ADDED
@@ -0,0 +1,63 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "format": {
3
+ "8-bit layers": "uint8 blocks of 32 with fp32 scale + uint8 zero point, dequantized in-graph (Cast/Sub/Mul, folded at load)",
4
+ "4-bit layers": "uint4 blocks of 32, MatMulNBits"
5
+ },
6
+ "group_size": 32,
7
+ "calib_pages": 256,
8
+ "eval_pages": 400,
9
+ "fp32_layers": [
10
+ "cnn.proj",
11
+ "page_in.net.0",
12
+ "page_in.net.3",
13
+ "blocks.0.rel",
14
+ "blocks.1.rel",
15
+ "blocks.2.rel",
16
+ "blocks.3.rel",
17
+ "blocks.4.rel",
18
+ "blocks.5.rel",
19
+ "blocks.6.rel",
20
+ "blocks.7.rel"
21
+ ],
22
+ "int4_8bit_edge_layers": [
23
+ "tok_in.net.0",
24
+ "tok_in.net.3",
25
+ "img_in.net.0",
26
+ "img_in.net.3",
27
+ "color_head.net.0",
28
+ "color_head.net.3",
29
+ "img_head.net.0",
30
+ "img_head.net.3"
31
+ ],
32
+ "variants": {
33
+ "fp32": {
34
+ "haiku_dE": 0.039349029815445344,
35
+ "native_dE": 0.08639045081450604,
36
+ "img_acc": 0.7966292134831461,
37
+ "size_mb": 63.7,
38
+ "cpu_ms_median": 11.42
39
+ },
40
+ "int8": {
41
+ "dE_vs_fp32": 0.00015194275738394936,
42
+ "dE_max": 0.02521813102066517,
43
+ "img_agree": 0.9998743718592965,
44
+ "haiku_dE": 0.03935358844697476,
45
+ "native_dE": 0.0863548206735868,
46
+ "img_acc": 0.7968164794007491,
47
+ "pack_parity_max_abs": 2.5033950805664062e-06,
48
+ "size_mb": 20.63,
49
+ "cpu_ms_median": 9.19
50
+ },
51
+ "int4": {
52
+ "dE_vs_fp32": 0.00036296586506068707,
53
+ "dE_max": 0.0709674283862114,
54
+ "img_agree": 0.9976130653266332,
55
+ "haiku_dE": 0.039340872559696434,
56
+ "native_dE": 0.08632776886341162,
57
+ "img_acc": 0.797003745318352,
58
+ "pack_parity_max_abs": 1.5497207641601562e-06,
59
+ "size_mb": 13.31,
60
+ "cpu_ms_median": 9.76
61
+ }
62
+ }
63
+ }
reports/train_log.json ADDED
@@ -0,0 +1,501 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "args": {
3
+ "data": "/home/user/.local/share/rtx-temporary-jobs/tj-nadir-0fe76575c0f3/payload/work/data/ds_final.pkl",
4
+ "out": "/job/export/run_m",
5
+ "d": 384,
6
+ "layers": 8,
7
+ "heads": 6,
8
+ "epochs": 30,
9
+ "bs": 48,
10
+ "lr": 0.0004,
11
+ "wd": 0.05,
12
+ "drop": 0.1,
13
+ "native_weight": 0.5,
14
+ "seed": 0,
15
+ "workers": 4
16
+ },
17
+ "params": 15895995,
18
+ "log": [
19
+ {
20
+ "epoch": 0,
21
+ "step": 468,
22
+ "secs": 81,
23
+ "loss": 0.3405,
24
+ "dE": 0.12132,
25
+ "con": 0.50822,
26
+ "img": 1.29224,
27
+ "s0_dE": 0.06306,
28
+ "s0_dL": 0.06024,
29
+ "s0_con": 0.04883,
30
+ "img_acc": 0.73416,
31
+ "s1_dE": 0.16366,
32
+ "s1_dL": 0.15909,
33
+ "s1_con": 0.11938
34
+ },
35
+ {
36
+ "epoch": 1,
37
+ "step": 936,
38
+ "secs": 148,
39
+ "loss": 0.18873,
40
+ "dE": 0.06161,
41
+ "con": 0.03771,
42
+ "img": 0.99138,
43
+ "s0_dE": 0.05432,
44
+ "s0_dL": 0.05155,
45
+ "s0_con": 0.03668,
46
+ "img_acc": 0.78603,
47
+ "s1_dE": 0.1523,
48
+ "s1_dL": 0.14838,
49
+ "s1_con": 0.08049
50
+ },
51
+ {
52
+ "epoch": 2,
53
+ "step": 1404,
54
+ "secs": 215,
55
+ "loss": 0.17375,
56
+ "dE": 0.05668,
57
+ "con": 0.03498,
58
+ "img": 0.90371,
59
+ "s0_dE": 0.04977,
60
+ "s0_dL": 0.04725,
61
+ "s0_con": 0.02243,
62
+ "img_acc": 0.78005,
63
+ "s1_dE": 0.14476,
64
+ "s1_dL": 0.14153,
65
+ "s1_con": 0.26806
66
+ },
67
+ {
68
+ "epoch": 3,
69
+ "step": 1872,
70
+ "secs": 281,
71
+ "loss": 0.16779,
72
+ "dE": 0.05488,
73
+ "con": 0.04139,
74
+ "img": 0.86775,
75
+ "s0_dE": 0.0506,
76
+ "s0_dL": 0.04808,
77
+ "s0_con": 0.0331,
78
+ "img_acc": 0.79611,
79
+ "s1_dE": 0.13697,
80
+ "s1_dL": 0.13372,
81
+ "s1_con": 0.1237
82
+ },
83
+ {
84
+ "epoch": 4,
85
+ "step": 2340,
86
+ "secs": 349,
87
+ "loss": 0.16278,
88
+ "dE": 0.05274,
89
+ "con": 0.0414,
90
+ "img": 0.85412,
91
+ "s0_dE": 0.04849,
92
+ "s0_dL": 0.04604,
93
+ "s0_con": 0.02439,
94
+ "img_acc": 0.78982,
95
+ "s1_dE": 0.13195,
96
+ "s1_dL": 0.12841,
97
+ "s1_con": 0.16648
98
+ },
99
+ {
100
+ "epoch": 5,
101
+ "step": 2808,
102
+ "secs": 416,
103
+ "loss": 0.15903,
104
+ "dE": 0.05166,
105
+ "con": 0.03833,
106
+ "img": 0.84086,
107
+ "s0_dE": 0.0472,
108
+ "s0_dL": 0.04477,
109
+ "s0_con": 0.03465,
110
+ "img_acc": 0.80467,
111
+ "s1_dE": 0.12882,
112
+ "s1_dL": 0.12538,
113
+ "s1_con": 0.12963
114
+ },
115
+ {
116
+ "epoch": 6,
117
+ "step": 3276,
118
+ "secs": 484,
119
+ "loss": 0.15729,
120
+ "dE": 0.05076,
121
+ "con": 0.04035,
122
+ "img": 0.83185,
123
+ "s0_dE": 0.0457,
124
+ "s0_dL": 0.04338,
125
+ "s0_con": 0.02999,
126
+ "img_acc": 0.80345,
127
+ "s1_dE": 0.15936,
128
+ "s1_dL": 0.15619,
129
+ "s1_con": 0.51491
130
+ },
131
+ {
132
+ "epoch": 7,
133
+ "step": 3744,
134
+ "secs": 552,
135
+ "loss": 0.15392,
136
+ "dE": 0.04978,
137
+ "con": 0.0403,
138
+ "img": 0.81841,
139
+ "s0_dE": 0.04612,
140
+ "s0_dL": 0.04395,
141
+ "s0_con": 0.01962,
142
+ "img_acc": 0.79899,
143
+ "s1_dE": 0.14854,
144
+ "s1_dL": 0.14578,
145
+ "s1_con": 0.42056
146
+ },
147
+ {
148
+ "epoch": 8,
149
+ "step": 4212,
150
+ "secs": 619,
151
+ "loss": 0.15244,
152
+ "dE": 0.04894,
153
+ "con": 0.04062,
154
+ "img": 0.81195,
155
+ "s0_dE": 0.04632,
156
+ "s0_dL": 0.04425,
157
+ "s0_con": 0.02413,
158
+ "img_acc": 0.79917,
159
+ "s1_dE": 0.12446,
160
+ "s1_dL": 0.12131,
161
+ "s1_con": 0.23072
162
+ },
163
+ {
164
+ "epoch": 9,
165
+ "step": 4680,
166
+ "secs": 687,
167
+ "loss": 0.14931,
168
+ "dE": 0.0481,
169
+ "con": 0.03785,
170
+ "img": 0.79837,
171
+ "s0_dE": 0.04452,
172
+ "s0_dL": 0.04238,
173
+ "s0_con": 0.02434,
174
+ "img_acc": 0.80094,
175
+ "s1_dE": 0.12096,
176
+ "s1_dL": 0.11779,
177
+ "s1_con": 0.21844
178
+ },
179
+ {
180
+ "epoch": 10,
181
+ "step": 5148,
182
+ "secs": 755,
183
+ "loss": 0.1484,
184
+ "dE": 0.04786,
185
+ "con": 0.03854,
186
+ "img": 0.79393,
187
+ "s0_dE": 0.04549,
188
+ "s0_dL": 0.04344,
189
+ "s0_con": 0.02666,
190
+ "img_acc": 0.79123,
191
+ "s1_dE": 0.11822,
192
+ "s1_dL": 0.11502,
193
+ "s1_con": 0.25328
194
+ },
195
+ {
196
+ "epoch": 11,
197
+ "step": 5616,
198
+ "secs": 822,
199
+ "loss": 0.14645,
200
+ "dE": 0.04703,
201
+ "con": 0.03983,
202
+ "img": 0.78304,
203
+ "s0_dE": 0.04345,
204
+ "s0_dL": 0.04139,
205
+ "s0_con": 0.01948,
206
+ "img_acc": 0.81053,
207
+ "s1_dE": 0.11486,
208
+ "s1_dL": 0.11153,
209
+ "s1_con": 0.20122
210
+ },
211
+ {
212
+ "epoch": 12,
213
+ "step": 6084,
214
+ "secs": 891,
215
+ "loss": 0.13803,
216
+ "dE": 0.04635,
217
+ "con": 0.03723,
218
+ "img": 0.76961,
219
+ "s0_dE": 0.04328,
220
+ "s0_dL": 0.04125,
221
+ "s0_con": 0.01831,
222
+ "img_acc": 0.8051,
223
+ "s1_dE": 0.11435,
224
+ "s1_dL": 0.11091,
225
+ "s1_con": 0.16656
226
+ },
227
+ {
228
+ "epoch": 13,
229
+ "step": 6552,
230
+ "secs": 958,
231
+ "loss": 0.13627,
232
+ "dE": 0.04589,
233
+ "con": 0.03889,
234
+ "img": 0.76454,
235
+ "s0_dE": 0.04247,
236
+ "s0_dL": 0.04047,
237
+ "s0_con": 0.01875,
238
+ "img_acc": 0.80161,
239
+ "s1_dE": 0.11794,
240
+ "s1_dL": 0.11469,
241
+ "s1_con": 0.23754
242
+ },
243
+ {
244
+ "epoch": 14,
245
+ "step": 7020,
246
+ "secs": 1025,
247
+ "loss": 0.13324,
248
+ "dE": 0.045,
249
+ "con": 0.0383,
250
+ "img": 0.75126,
251
+ "s0_dE": 0.04194,
252
+ "s0_dL": 0.03998,
253
+ "s0_con": 0.02424,
254
+ "img_acc": 0.80332,
255
+ "s1_dE": 0.11805,
256
+ "s1_dL": 0.11474,
257
+ "s1_con": 0.22873
258
+ },
259
+ {
260
+ "epoch": 15,
261
+ "step": 7488,
262
+ "secs": 1093,
263
+ "loss": 0.13267,
264
+ "dE": 0.04502,
265
+ "con": 0.0394,
266
+ "img": 0.74724,
267
+ "s0_dE": 0.04183,
268
+ "s0_dL": 0.03983,
269
+ "s0_con": 0.01954,
270
+ "img_acc": 0.81353,
271
+ "s1_dE": 0.11306,
272
+ "s1_dL": 0.10954,
273
+ "s1_con": 0.20409
274
+ },
275
+ {
276
+ "epoch": 16,
277
+ "step": 7956,
278
+ "secs": 1161,
279
+ "loss": 0.1305,
280
+ "dE": 0.04465,
281
+ "con": 0.03933,
282
+ "img": 0.72731,
283
+ "s0_dE": 0.04195,
284
+ "s0_dL": 0.04006,
285
+ "s0_con": 0.02284,
286
+ "img_acc": 0.80485,
287
+ "s1_dE": 0.11437,
288
+ "s1_dL": 0.11122,
289
+ "s1_con": 0.27215
290
+ },
291
+ {
292
+ "epoch": 17,
293
+ "step": 8424,
294
+ "secs": 1229,
295
+ "loss": 0.12864,
296
+ "dE": 0.04415,
297
+ "con": 0.03796,
298
+ "img": 0.71419,
299
+ "s0_dE": 0.04177,
300
+ "s0_dL": 0.0399,
301
+ "s0_con": 0.0207,
302
+ "img_acc": 0.79611,
303
+ "s1_dE": 0.10924,
304
+ "s1_dL": 0.10592,
305
+ "s1_con": 0.22063
306
+ },
307
+ {
308
+ "epoch": 18,
309
+ "step": 8892,
310
+ "secs": 1296,
311
+ "loss": 0.12691,
312
+ "dE": 0.04379,
313
+ "con": 0.03646,
314
+ "img": 0.69736,
315
+ "s0_dE": 0.04175,
316
+ "s0_dL": 0.03993,
317
+ "s0_con": 0.02416,
318
+ "img_acc": 0.801,
319
+ "s1_dE": 0.10774,
320
+ "s1_dL": 0.10449,
321
+ "s1_con": 0.24435
322
+ },
323
+ {
324
+ "epoch": 19,
325
+ "step": 9360,
326
+ "secs": 1364,
327
+ "loss": 0.12499,
328
+ "dE": 0.04319,
329
+ "con": 0.036,
330
+ "img": 0.68816,
331
+ "s0_dE": 0.0413,
332
+ "s0_dL": 0.03945,
333
+ "s0_con": 0.02047,
334
+ "img_acc": 0.79813,
335
+ "s1_dE": 0.11372,
336
+ "s1_dL": 0.11037,
337
+ "s1_con": 0.20421
338
+ },
339
+ {
340
+ "epoch": 20,
341
+ "step": 9828,
342
+ "secs": 1432,
343
+ "loss": 0.12434,
344
+ "dE": 0.04334,
345
+ "con": 0.03581,
346
+ "img": 0.67546,
347
+ "s0_dE": 0.0409,
348
+ "s0_dL": 0.03905,
349
+ "s0_con": 0.01907,
350
+ "img_acc": 0.79483,
351
+ "s1_dE": 0.10692,
352
+ "s1_dL": 0.1036,
353
+ "s1_con": 0.22495
354
+ },
355
+ {
356
+ "epoch": 21,
357
+ "step": 10296,
358
+ "secs": 1499,
359
+ "loss": 0.12248,
360
+ "dE": 0.04287,
361
+ "con": 0.03581,
362
+ "img": 0.65995,
363
+ "s0_dE": 0.04105,
364
+ "s0_dL": 0.0392,
365
+ "s0_con": 0.02163,
366
+ "img_acc": 0.7999,
367
+ "s1_dE": 0.10461,
368
+ "s1_dL": 0.10127,
369
+ "s1_con": 0.18494
370
+ },
371
+ {
372
+ "epoch": 22,
373
+ "step": 10764,
374
+ "secs": 1567,
375
+ "loss": 0.12033,
376
+ "dE": 0.04222,
377
+ "con": 0.03395,
378
+ "img": 0.64253,
379
+ "s0_dE": 0.04075,
380
+ "s0_dL": 0.03893,
381
+ "s0_con": 0.02266,
382
+ "img_acc": 0.7999,
383
+ "s1_dE": 0.10178,
384
+ "s1_dL": 0.09855,
385
+ "s1_con": 0.23005
386
+ },
387
+ {
388
+ "epoch": 23,
389
+ "step": 11232,
390
+ "secs": 1635,
391
+ "loss": 0.11905,
392
+ "dE": 0.04192,
393
+ "con": 0.03691,
394
+ "img": 0.63201,
395
+ "s0_dE": 0.04024,
396
+ "s0_dL": 0.03843,
397
+ "s0_con": 0.01664,
398
+ "img_acc": 0.79049,
399
+ "s1_dE": 0.10458,
400
+ "s1_dL": 0.10125,
401
+ "s1_con": 0.1973
402
+ },
403
+ {
404
+ "epoch": 24,
405
+ "step": 11700,
406
+ "secs": 1703,
407
+ "loss": 0.11767,
408
+ "dE": 0.04169,
409
+ "con": 0.03462,
410
+ "img": 0.61761,
411
+ "s0_dE": 0.04043,
412
+ "s0_dL": 0.03861,
413
+ "s0_con": 0.02038,
414
+ "img_acc": 0.79379,
415
+ "s1_dE": 0.1025,
416
+ "s1_dL": 0.09926,
417
+ "s1_con": 0.2135
418
+ },
419
+ {
420
+ "epoch": 25,
421
+ "step": 12168,
422
+ "secs": 1770,
423
+ "loss": 0.11655,
424
+ "dE": 0.04131,
425
+ "con": 0.03564,
426
+ "img": 0.60858,
427
+ "s0_dE": 0.04068,
428
+ "s0_dL": 0.0389,
429
+ "s0_con": 0.01981,
430
+ "img_acc": 0.79715,
431
+ "s1_dE": 0.10372,
432
+ "s1_dL": 0.10038,
433
+ "s1_con": 0.20402
434
+ },
435
+ {
436
+ "epoch": 26,
437
+ "step": 12636,
438
+ "secs": 1838,
439
+ "loss": 0.11502,
440
+ "dE": 0.04106,
441
+ "con": 0.03415,
442
+ "img": 0.59825,
443
+ "s0_dE": 0.04032,
444
+ "s0_dL": 0.03852,
445
+ "s0_con": 0.01906,
446
+ "img_acc": 0.79691,
447
+ "s1_dE": 0.10293,
448
+ "s1_dL": 0.09969,
449
+ "s1_con": 0.22009
450
+ },
451
+ {
452
+ "epoch": 27,
453
+ "step": 13104,
454
+ "secs": 1906,
455
+ "loss": 0.11521,
456
+ "dE": 0.04109,
457
+ "con": 0.03603,
458
+ "img": 0.59315,
459
+ "s0_dE": 0.0404,
460
+ "s0_dL": 0.03861,
461
+ "s0_con": 0.02053,
462
+ "img_acc": 0.79721,
463
+ "s1_dE": 0.10302,
464
+ "s1_dL": 0.09978,
465
+ "s1_con": 0.2078
466
+ },
467
+ {
468
+ "epoch": 28,
469
+ "step": 13572,
470
+ "secs": 1973,
471
+ "loss": 0.11402,
472
+ "dE": 0.04082,
473
+ "con": 0.03683,
474
+ "img": 0.58722,
475
+ "s0_dE": 0.04034,
476
+ "s0_dL": 0.03856,
477
+ "s0_con": 0.01831,
478
+ "img_acc": 0.79098,
479
+ "s1_dE": 0.10383,
480
+ "s1_dL": 0.10053,
481
+ "s1_con": 0.181
482
+ },
483
+ {
484
+ "epoch": 29,
485
+ "step": 14040,
486
+ "secs": 2041,
487
+ "loss": 0.11321,
488
+ "dE": 0.04059,
489
+ "con": 0.03495,
490
+ "img": 0.5803,
491
+ "s0_dE": 0.04018,
492
+ "s0_dL": 0.03841,
493
+ "s0_con": 0.01978,
494
+ "img_acc": 0.79416,
495
+ "s1_dE": 0.10315,
496
+ "s1_dL": 0.09986,
497
+ "s1_con": 0.19971
498
+ }
499
+ ],
500
+ "best_score": 0.07564117624840601
501
+ }
runtime/DARKREADER-LICENSE ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Dark Reader Ltd.
4
+
5
+ All rights reserved.
6
+
7
+ Permission is hereby granted, free of charge, to any person obtaining a copy
8
+ of this software and associated documentation files (the "Software"), to deal
9
+ in the Software without restriction, including without limitation the rights
10
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
11
+ copies of the Software, and to permit persons to whom the Software is
12
+ furnished to do so, subject to the following conditions:
13
+
14
+ The above copyright notice and this permission notice shall be included in all
15
+ copies or substantial portions of the Software.
16
+
17
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
18
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
19
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
20
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
21
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
22
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
23
+ SOFTWARE.
runtime/darkreader-nadir.js ADDED
The diff for this file is too large to render. See raw diff
 
runtime/extract.js ADDED
@@ -0,0 +1,825 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ // Nadir page analyzer. Shared by the training crawler (Playwright) and the browser extension.
2
+ // Defines globalThis.__nadirExtract(opts) -> Promise<PageDescription>.
3
+ // Everything the model sees is computed here, so training and inference features match exactly.
4
+ (function () {
5
+ 'use strict';
6
+ if (globalThis.__nadirExtract) return;
7
+
8
+ const ROLES = ['bg', 'text', 'border', 'shadow', 'gradient', 'fill', 'stroke'];
9
+ const R = Object.fromEntries(ROLES.map((r, i) => [r, i]));
10
+ const SKIP_TAGS = new Set(['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEMPLATE', 'HEAD', 'META', 'LINK', 'TITLE', 'BR', 'WBR', 'SOURCE', 'TRACK', 'PARAM']);
11
+ const INTERACTIVE = new Set(['A', 'BUTTON', 'INPUT', 'SELECT', 'TEXTAREA', 'SUMMARY', 'LABEL']);
12
+ const FORM = new Set(['INPUT', 'SELECT', 'TEXTAREA']);
13
+ const HEADINGS = new Set(['H1', 'H2', 'H3', 'H4', 'H5', 'H6']);
14
+ const CODE = new Set(['CODE', 'PRE', 'KBD', 'SAMP']);
15
+ const SVG_SHAPES = new Set(['path', 'circle', 'rect', 'ellipse', 'line', 'polyline', 'polygon', 'text', 'tspan', 'use', 'g', 'svg']);
16
+
17
+ // ---------- color parsing ----------
18
+ const colorCache = new Map();
19
+ let cctx = null;
20
+ const BAD_COLOR = /var\(|env\(|attr\(|calc\(|^(inherit|initial|unset|revert|revert-layer|currentcolor|auto|none|transparent|-webkit-[\w-]+|-moz-[\w-]+)$/i;
21
+ function splitTop(s) { // split on top-level commas
22
+ const out = []; let depth = 0, cur = '';
23
+ for (const ch of s) { if (ch === '(') depth++; else if (ch === ')') depth--; if (ch === ',' && depth === 0) { out.push(cur); cur = ''; } else cur += ch; }
24
+ out.push(cur); return out.map(x => x.trim());
25
+ }
26
+ function canvasParse(str) {
27
+ const ld = str.match(/^light-dark\((.*)\)$/i);
28
+ if (ld) { const parts = splitTop(ld[1]); return parts.length === 2 ? parseColor(parts[0]) : null; }
29
+ if (BAD_COLOR.test(str)) return null;
30
+ if (typeof CSS !== 'undefined' && CSS.supports && !CSS.supports('color', str)) return null;
31
+ if (!cctx) {
32
+ const c = (typeof OffscreenCanvas !== 'undefined') ? new OffscreenCanvas(1, 1) : Object.assign(document.createElement('canvas'), {width: 1, height: 1});
33
+ cctx = c.getContext('2d', {willReadFrequently: true});
34
+ }
35
+ cctx.fillStyle = '#010203';
36
+ const before = cctx.fillStyle;
37
+ cctx.fillStyle = str;
38
+ if (cctx.fillStyle === before && !/^#?010203/i.test(str)) return null; // canvas rejected it
39
+ cctx.clearRect(0, 0, 1, 1);
40
+ cctx.fillRect(0, 0, 1, 1);
41
+ const d = cctx.getImageData(0, 0, 1, 1).data;
42
+ return [d[0], d[1], d[2], d[3] / 255];
43
+ }
44
+ function parseColor(str) {
45
+ if (!str) return null;
46
+ let c = colorCache.get(str);
47
+ if (c !== undefined) return c;
48
+ c = null;
49
+ const s = str.trim();
50
+ if (s === 'transparent' || s === 'none' || s === '' || s.startsWith('url(') || s === 'currentcolor') c = null;
51
+ else {
52
+ const m = s.match(/^rgba?\(\s*([\d.]+)[\s,]+([\d.]+)[\s,]+([\d.]+)(?:\s*[,/]\s*([\d.]+%?))?\s*\)$/);
53
+ if (m) {
54
+ let a = m[4] === undefined ? 1 : (m[4].endsWith('%') ? parseFloat(m[4]) / 100 : parseFloat(m[4]));
55
+ c = [Math.round(+m[1]), Math.round(+m[2]), Math.round(+m[3]), a];
56
+ } else {
57
+ try { c = canvasParse(s); } catch (e) { c = null; }
58
+ }
59
+ if (c && c[3] <= 0.004) c = null;
60
+ }
61
+ colorCache.set(str, c);
62
+ return c;
63
+ }
64
+ // all colors inside a composite value (box-shadow, gradients)
65
+ const COLOR_RE = /(rgba?\([^)]*\)|hsla?\([^)]*\)|oklch\([^)]*\)|oklab\([^)]*\)|lab\([^)]*\)|lch\([^)]*\)|color\([^)]*\)|#[0-9a-fA-F]{3,8}\b)/g;
66
+ function colorsIn(str) {
67
+ if (!str || str === 'none') return [];
68
+ const out = [];
69
+ for (const m of str.matchAll(COLOR_RE)) { const c = parseColor(m[1]); if (c) out.push(c); }
70
+ return out;
71
+ }
72
+ const hex2 = (n) => n.toString(16).padStart(2, '0');
73
+ const keyOf = (c) => hex2(c[0]) + hex2(c[1]) + hex2(c[2]) + hex2(Math.round(c[3] * 255));
74
+
75
+ // sRGB -> OKLab
76
+ function lin(u) { u /= 255; return u <= 0.04045 ? u / 12.92 : Math.pow((u + 0.055) / 1.055, 2.4); }
77
+ function oklab(c) {
78
+ const r = lin(c[0]), g = lin(c[1]), b = lin(c[2]);
79
+ const l = Math.cbrt(0.4122214708 * r + 0.5363325363 * g + 0.0514459929 * b);
80
+ const m = Math.cbrt(0.2119034982 * r + 0.6806995451 * g + 0.1073969566 * b);
81
+ const s = Math.cbrt(0.0883024619 * r + 0.2817188376 * g + 0.6299787005 * b);
82
+ return [0.2104542553 * l + 0.7936177850 * m - 0.0040720468 * s,
83
+ 1.9779984951 * l - 2.4285922050 * m + 0.4505937099 * s,
84
+ 0.0259040371 * l + 0.7827717662 * m - 0.8086757660 * s];
85
+ }
86
+ function relLum(c) { return 0.2126 * lin(c[0]) + 0.7152 * lin(c[1]) + 0.0722 * lin(c[2]); }
87
+ function over(top, bottom) { // composite rgba over opaque rgb
88
+ const a = top[3];
89
+ return [top[0] * a + bottom[0] * (1 - a), top[1] * a + bottom[1] * (1 - a), top[2] * a + bottom[2] * (1 - a), 1];
90
+ }
91
+
92
+ // ---------- dark-mode support scan ----------
93
+ const DARK_SEL_RE = /(?:^|[\s,>+~(])((?:html|body|:root)?(?:\.(?:[\w-]*-)?dark(?:-[\w-]+)?|\.(?:[\w-]*-)?night[\w-]*|\[data-[\w-]*(?:theme|mode|scheme|color-mode|appearance)[\w-]*(?:[~|^*$]?=\s*["']?[\w-]*dark[\w-]*["']?)\s*(?:i\s*)?\]|\[(?:theme|mode)\s*=\s*["']?dark["']?\]|\[dark(?:-theme|-mode)?\]))/i;
94
+ function scanSheets(limitRules) {
95
+ const res = {sheets: 0, blocked: 0, rules: 0, darkMediaRules: 0, darkMediaColorDecls: 0, lightMediaRules: 0,
96
+ darkSelectors: {}, ruleColors: [], hoverColors: 0};
97
+ let budget = limitRules;
98
+ const seenRuleColors = new Map();
99
+ const pushRuleColor = (role, c, state) => {
100
+ const k = role + ':' + keyOf(c);
101
+ const e = seenRuleColors.get(k);
102
+ if (e) { e.n++; if (state) e.state++; return; }
103
+ seenRuleColors.set(k, {role, c, n: 1, state: state ? 1 : 0});
104
+ };
105
+ function walk(rules, inDarkMedia, inLightMedia) {
106
+ for (const rule of rules) {
107
+ if (budget-- <= 0) return;
108
+ if (rule.cssRules && rule.media) {
109
+ const mt = rule.media.mediaText || '';
110
+ const dark = /prefers-color-scheme\s*:\s*dark/i.test(mt);
111
+ const light = /prefers-color-scheme\s*:\s*light/i.test(mt);
112
+ if (dark) res.darkMediaRules++;
113
+ if (light) res.lightMediaRules++;
114
+ walk(rule.cssRules, inDarkMedia || dark, inLightMedia || light);
115
+ continue;
116
+ }
117
+ if (rule.cssRules && !rule.style) { walk(rule.cssRules, inDarkMedia, inLightMedia); continue; }
118
+ if (!rule.style) continue;
119
+ res.rules++;
120
+ const st = rule.style;
121
+ let colorDecls = 0;
122
+ const sel = rule.selectorText || '';
123
+ const state = /:(hover|focus|active|focus-visible|focus-within|checked)/.test(sel);
124
+ for (let i = 0; i < st.length; i++) {
125
+ const p = st[i];
126
+ const v = st.getPropertyValue(p);
127
+ if (!v) continue;
128
+ let role = -1;
129
+ if (p === 'background-color' || p === 'background') role = R.bg;
130
+ else if (p === 'color' || p === '-webkit-text-fill-color') role = R.text;
131
+ else if (p.startsWith('border') && p.endsWith('color') || p === 'outline-color') role = R.border;
132
+ else if (p === 'fill') role = R.fill;
133
+ else if (p === 'stroke') role = R.stroke;
134
+ else if (p === 'box-shadow' || p === 'text-shadow') role = R.shadow;
135
+ else if (p === 'background-image') role = R.gradient;
136
+ else if (p.startsWith('--')) { if (COLOR_RE.test(v)) colorDecls++; COLOR_RE.lastIndex = 0; continue; }
137
+ if (role < 0) continue;
138
+ colorDecls++;
139
+ if (inDarkMedia || inLightMedia) continue; // light page colors only
140
+ const cs = (role === R.shadow || role === R.gradient || p === 'background') ? colorsIn(v) : [parseColor(v)].filter(Boolean);
141
+ for (const c of cs) { pushRuleColor(role === R.gradient && p === 'background' ? R.bg : role, c, state); if (state) res.hoverColors++; }
142
+ }
143
+ if (inDarkMedia) res.darkMediaColorDecls += colorDecls;
144
+ if (colorDecls && sel) {
145
+ const m = sel.match(DARK_SEL_RE);
146
+ if (m) {
147
+ const k = m[1].trim();
148
+ res.darkSelectors[k] = (res.darkSelectors[k] || 0) + colorDecls;
149
+ }
150
+ }
151
+ }
152
+ }
153
+ const sheets = [...document.styleSheets, ...(document.adoptedStyleSheets || []), ...(scanSheets.extra || [])];
154
+ for (const sh of sheets) {
155
+ if (scanSheets.filter && !scanSheets.filter(sh)) continue;
156
+ res.sheets++;
157
+ let rules;
158
+ try { rules = sh.cssRules; } catch (e) { res.blocked++; continue; }
159
+ if (!rules) { res.blocked++; continue; }
160
+ walk(rules, sh.media && /prefers-color-scheme\s*:\s*dark/i.test(sh.media.mediaText || ''), false);
161
+ }
162
+ res.ruleColors = [...seenRuleColors.values()];
163
+ return res;
164
+ }
165
+
166
+ function schemeInfo(sheetScan) {
167
+ const meta = document.querySelector('meta[name="color-scheme"]');
168
+ const rootCS = getComputedStyle(document.documentElement).colorScheme || '';
169
+ const darkLinks = [...document.querySelectorAll('link[rel~="stylesheet"][media]')].filter(l => /prefers-color-scheme\s*:\s*dark/i.test(l.media)).length;
170
+ const darkSources = [...document.querySelectorAll('picture source[media]')].filter(s => /prefers-color-scheme\s*:\s*dark/i.test(s.media)).length;
171
+ const html = document.documentElement, body = document.body;
172
+ const attrs = {};
173
+ for (const el of [html, body]) {
174
+ if (!el) continue;
175
+ for (const a of el.attributes) if (/theme|mode|scheme|appearance|dark|light/i.test(a.name + '=' + a.value) && a.name !== 'style') attrs[el.tagName.toLowerCase() + '@' + a.name] = a.value.slice(0, 80);
176
+ }
177
+ const sel = Object.entries(sheetScan.darkSelectors).sort((a, b) => b[1] - a[1]).slice(0, 12);
178
+ return {
179
+ metaColorScheme: meta ? meta.content : null,
180
+ rootColorScheme: rootCS,
181
+ darkMediaRules: sheetScan.darkMediaRules,
182
+ darkMediaColorDecls: sheetScan.darkMediaColorDecls,
183
+ lightMediaRules: sheetScan.lightMediaRules,
184
+ darkLinks, darkSources,
185
+ darkSelectors: sel,
186
+ rootAttrs: attrs,
187
+ rootClass: (html.className && typeof html.className === 'string' ? html.className : '').slice(0, 200),
188
+ bodyClass: (body && typeof body.className === 'string' ? body.className : '').slice(0, 200),
189
+ sheets: sheetScan.sheets, blockedSheets: sheetScan.blocked, rules: sheetScan.rules,
190
+ };
191
+ }
192
+
193
+ // ---------- ads ----------
194
+ // Ad slots are ignored entirely, so a page looks the same to the model with or without an ad blocker.
195
+ const AD_TOKEN = /^(ad|ads|adv|advert|adverts|advertisement|advertising|adsbygoogle|adslot|adunit|adbox|adblock|adbanner|ad[-_](slot|unit|container|wrapper|banner|box|block|placeholder|label|space|zone|frame|holder|area|leaderboard|rail|sidebar|top|bottom|inline|native|mpu|billboard|skyscraper|sticky)|(top|side|bottom|inline|sticky|native|header|footer|right|left)[-_]?ads?|banner[-_]?ads?|gpt[-_]ad.*|dfp[-_]?ad.*|google[-_]ads?.*|taboola.*|trc_.*|outbrain.*|ob[-_]widget.*|sponsored([-_].*)?|promoted[-_]?(content|post|ad)s?|mgid.*|revcontent.*|teads.*|criteo.*|amzn[-_]?ad.*|adthrive.*|ezoic.*|mediavine.*)$/i;
196
+ const AD_ID = /^(google_ads|div-gpt-ad|gpt-ad|ad[-_]|ads[-_]|taboola|outbrain|adthrive|ezoic-pub-ad|mgid|revcontent)/i;
197
+ const AD_SRC = /(doubleclick\.net|googlesyndication\.com|googleadservices\.com|adservice\.google|amazon-adsystem\.com|adnxs\.com|taboola\.com|outbrain\.com|criteo\.|pubmatic\.com|rubiconproject\.com|openx\.net|adsrvr\.org|moatads\.com|mgid\.com|revcontent\.com|teads\.tv|media\.net\/|ezoic|adthrive|mediavine)/i;
198
+ function isAd(el) {
199
+ const tag = el.tagName;
200
+ if (tag === 'INS' && /adsbygoogle/.test(el.className || '')) return true;
201
+ if ((tag === 'IFRAME' || tag === 'IMG' || tag === 'A') && AD_SRC.test(el.getAttribute('src') || el.getAttribute('href') || '')) return true;
202
+ if (el.hasAttribute('data-ad-slot') || el.hasAttribute('data-google-query-id') || el.hasAttribute('data-ad-unit') || el.hasAttribute('data-adunit') || el.hasAttribute('data-ad-client')) return true;
203
+ const id = el.id;
204
+ if (id && typeof id === 'string' && AD_ID.test(id)) return true;
205
+ const cls = typeof el.className === 'string' ? el.className : '';
206
+ if (cls) { for (const c of cls.split(/\s+/)) if (c && AD_TOKEN.test(c)) return true; }
207
+ const al = el.getAttribute('aria-label');
208
+ if (al && /^(advertisement|ad|sponsored|anzeige|publicidad|publicité|pubblicità|реклама|広告)$/i.test(al.trim())) return true;
209
+ return false;
210
+ }
211
+
212
+ // ---------- element walk ----------
213
+ function describe(el) {
214
+ let s = el.tagName.toLowerCase();
215
+ if (el.id && typeof el.id === 'string') s += '#' + el.id.slice(0, 24);
216
+ const cls = typeof el.className === 'string' ? el.className : (el.className && el.className.baseVal) || '';
217
+ if (cls) s += '.' + cls.trim().split(/\s+/).slice(0, 3).join('.').slice(0, 48);
218
+ return s;
219
+ }
220
+ function ownText(el) {
221
+ let t = '';
222
+ for (const n of el.childNodes) if (n.nodeType === 3) t += n.nodeValue;
223
+ t = t.replace(/\s+/g, ' ').trim();
224
+ if (!t && FORM.has(el.tagName)) {
225
+ const v = el.value || el.placeholder || '';
226
+ if (el.type !== 'hidden' && el.type !== 'checkbox' && el.type !== 'radio') t = String(v).slice(0, 60) || ' ';
227
+ }
228
+ return t;
229
+ }
230
+
231
+ async function extract(opts) {
232
+ opts = Object.assign({maxElements: 5000, maxRules: 40000, maxTokens: 384, maxImages: 48, describe: false,
233
+ thumbs: false, thumbSize: 32, bigThumbSize: 0, fetchImage: null, pseudo: true, time: true}, opts || {});
234
+ const t0 = performance.now();
235
+ const vw = innerWidth, vh = innerHeight;
236
+ const sx = scrollX, sy = scrollY;
237
+ const docEl = document.documentElement;
238
+ const docH = Math.max(docEl.scrollHeight, document.body ? document.body.scrollHeight : 0, vh);
239
+ const docW = Math.max(vw, Math.min(docEl.scrollWidth, vw * 1.5));
240
+ const pageArea = docW * Math.min(docH, vh * 8);
241
+ scanSheets.filter = opts.sheetFilter || null;
242
+ scanSheets.extra = opts.extraSheets || [];
243
+ const sheetScan = scanSheets(opts.maxRules);
244
+
245
+ // canvas default background (white unless the page opts into dark color-scheme)
246
+ const rootCS = getComputedStyle(docEl);
247
+ const canvasDark = /dark/.test(rootCS.colorScheme) && !/light/.test(rootCS.colorScheme);
248
+ const canvasColor = canvasDark ? [18, 18, 18, 1] : [255, 255, 255, 1];
249
+
250
+ const tokens = new Map(); // key role:hex -> token
251
+ function tok(role, c) {
252
+ const k = role + ':' + keyOf(c);
253
+ let t = tokens.get(k);
254
+ if (!t) {
255
+ t = {k, role, c, i: -1, area: 0, area1: 0, n: 0, nVis: 0, chars: 0, fsSum: 0, fsMax: 0, bold: 0, root: 0,
256
+ inter: 0, link: 0, head: 0, code: 0, form: 0, header: 0, footer: 0, nav: 0, under: [0, 0, 0, 0], over: [0, 0, 0, 0],
257
+ contrastSum: 0, contrastW: 0, rules: 0, state: 0, depthSum: 0, ex: [], xs: 0, ys: 0, ws: 0, hs: 0, imgUnder: 0, fixed: 0, media: 0};
258
+ tokens.set(k, t);
259
+ }
260
+ return t;
261
+ }
262
+ const rel = new Map(); // "type|a|b" -> weight
263
+ function addRel(type, a, b, w) {
264
+ if (!a || !b || a === b) return;
265
+ const k = type + '|' + a.k + '|' + b.k;
266
+ rel.set(k, (rel.get(k) || 0) + w);
267
+ }
268
+
269
+ // walk DOM (including open shadow roots) in document order
270
+ let adsSkipped = 0;
271
+ const els = [];
272
+ const stack = document.body ? [document.body] : [];
273
+ const deferred = [];
274
+ const parentOf = new Map();
275
+ while ((stack.length || deferred.length) && els.length < opts.maxElements) {
276
+ const el = stack.length ? stack.pop() : deferred.shift();
277
+ els.push(el);
278
+ if (el !== document.body && el.getClientRects().length === 0 && !(el instanceof SVGElement)) {
279
+ // hidden subtree (display:none etc.): visit later, only if budget remains
280
+ const kids = [...el.children].filter(k => !SKIP_TAGS.has(k.tagName) && !isAd(k));
281
+ for (const k of kids) { parentOf.set(k, el); deferred.push(k); }
282
+ continue;
283
+ }
284
+ const kids = [];
285
+ const root = el.shadowRoot || null;
286
+ if (root) for (const k of root.children) kids.push(k);
287
+ for (const k of el.children) kids.push(k);
288
+ for (let i = kids.length - 1; i >= 0; i--) {
289
+ const k = kids[i];
290
+ if (SKIP_TAGS.has(k.tagName)) continue;
291
+ if (isAd(k)) { adsSkipped++; continue; }
292
+ parentOf.set(k, el);
293
+ stack.push(k);
294
+ }
295
+ }
296
+
297
+ // media (video players, embeds, canvases, large images): overlays on these must keep their colours
298
+ const media = [];
299
+ for (const m of document.querySelectorAll('video, iframe, canvas, embed, object, img')) {
300
+ if (m.tagName === 'IFRAME' && AD_SRC.test(m.src || '')) continue;
301
+ const r = m.getBoundingClientRect();
302
+ if (r.width * r.height < (m.tagName === 'IMG' ? 40000 : 6000)) continue;
303
+ media.push({m, x: r.left + sx, y: r.top + sy, w: r.width, h: r.height});
304
+ if (media.length > 64) break;
305
+ }
306
+ const onMedia = (el, x, y, w, h) => {
307
+ if (!media.length || w * h <= 0) return 0;
308
+ for (const md of media) {
309
+ if (md.m === el) continue;
310
+ if (el.contains(md.m) && w * h < md.w * md.h * 2.5) return 1; // player/media container
311
+ const ix = Math.max(0, Math.min(x + w, md.x + md.w) - Math.max(x, md.x));
312
+ const iy = Math.max(0, Math.min(y + h, md.y + md.h) - Math.max(y, md.y));
313
+ if (ix * iy >= 0.6 * w * h) return 1; // drawn on top of media
314
+ }
315
+ return 0;
316
+ };
317
+ // per-element info
318
+ const info = new Map();
319
+ const htmlBg = parseColor(rootCS.backgroundColor);
320
+ const bodyCS = document.body ? getComputedStyle(document.body) : null;
321
+ const bodyBg = bodyCS ? parseColor(bodyCS.backgroundColor) : null;
322
+ // effective root surface: body bg propagates to canvas if html has none
323
+ let rootSurface = htmlBg ? over(htmlBg, canvasColor) : (bodyBg ? over(bodyBg, canvasColor) : canvasColor);
324
+ const rootTok = tok(R.bg, htmlBg || bodyBg || canvasColor);
325
+ rootTok.root = 1;
326
+ rootTok.n++; rootTok.nVis++;
327
+ const rootRec = {surf: rootSurface, surfTok: rootTok, imgBg: false, region: 0, depth: 0, fixed: false};
328
+ const images = [];
329
+ const svgRoots = [];
330
+ let visibleCount = 0;
331
+
332
+ const ec = opts.elemColors ? [] : null;
333
+ for (const el of els) {
334
+ const cs = getComputedStyle(el);
335
+ const ecRow = ec ? [el.tagName, 0, null, null, null, null, null, null, null] : null;
336
+ if (ec) ec.push(ecRow);
337
+ const parent = parentOf.get(el);
338
+ const p = parent ? info.get(parent) : rootRec;
339
+ const pr = p || rootRec;
340
+ if (cs.display === 'none') { info.set(el, Object.assign({}, pr, {hidden: true})); }
341
+ const hidden = cs.display === 'none' || (pr && pr.hidden) || cs.visibility === 'hidden' || cs.opacity === '0';
342
+ const rect = el.getBoundingClientRect();
343
+ const x = rect.left + sx, y = rect.top + sy;
344
+ const w = Math.max(0, Math.min(rect.width, docW - Math.max(0, x)));
345
+ const h = Math.max(0, rect.height);
346
+ const vis = !hidden && w > 1 && h > 1 && x + w > 0 && y + h > 0;
347
+ const area = vis ? Math.min(w * h, docW * docH) : 0;
348
+ const area1 = vis ? Math.max(0, Math.min(x + w, vw) - Math.max(x, 0)) * Math.max(0, Math.min(y + h, vh) - Math.max(y, 0)) : 0;
349
+ if (vis) visibleCount++;
350
+ if (ecRow) ecRow[1] = Math.round(area);
351
+ const tag = el.tagName;
352
+ const isSvg = el instanceof SVGElement;
353
+ // regions
354
+ let region = pr.region;
355
+ if (tag === 'HEADER' || el.getAttribute && el.getAttribute('role') === 'banner') region |= 1;
356
+ if (tag === 'NAV' || el.getAttribute && el.getAttribute('role') === 'navigation') region |= 2;
357
+ if (tag === 'FOOTER' || el.getAttribute && el.getAttribute('role') === 'contentinfo') region |= 4;
358
+ const fixed = pr.fixed || cs.position === 'fixed' || cs.position === 'sticky';
359
+ const om = vis ? onMedia(el, x, y, w, h) : 0;
360
+ const rec = {surf: pr.surf, surfTok: pr.surfTok, imgBg: pr.imgBg, region, depth: pr.depth + 1, hidden: hidden, fixed, om,
361
+ inLink: pr.inLink || tag === 'A', link: tag === 'A' ? el : pr.link};
362
+ info.set(el, rec);
363
+ const flags = (t) => {
364
+ t.n++;
365
+ if (vis) { t.nVis++; t.depthSum += rec.depth; }
366
+ if (INTERACTIVE.has(tag) || el.getAttribute && el.getAttribute('role') === 'button') t.inter++;
367
+ if (rec.inLink) t.link++;
368
+ if (HEADINGS.has(tag)) t.head++;
369
+ if (CODE.has(tag)) t.code++;
370
+ if (FORM.has(tag)) t.form++;
371
+ if (region & 1) t.header++;
372
+ if (region & 2) t.nav++;
373
+ if (region & 4) t.footer++;
374
+ if (fixed) t.fixed++;
375
+ if (om) t.media++;
376
+ if (opts.describe && t.ex.length < 3 && vis) {
377
+ const txt = ownText(el);
378
+ t.ex.push(describe(el) + (txt ? ' "' + txt.slice(0, 40) + '"' : ''));
379
+ }
380
+ };
381
+
382
+ // background
383
+ const bg = parseColor(cs.backgroundColor);
384
+ const bgImg = cs.backgroundImage;
385
+ if (bg) {
386
+ const t = tok(R.bg, bg);
387
+ if (ecRow) ecRow[2] = keyOf(bg);
388
+ flags(t);
389
+ t.area += area; t.area1 += area1;
390
+ if (vis) { t.xs += x * area; t.ys += y * area; t.ws += w * area; t.hs += h * area; }
391
+ const under = pr.surf;
392
+ if (vis) { addRel('bb', t, pr.surfTok, area); }
393
+ const ul = oklab(under);
394
+ t.under[0] += ul[0] * (area + 1); t.under[1] += ul[1] * (area + 1); t.under[2] += ul[2] * (area + 1); t.under[3] += area + 1;
395
+ if (pr.imgBg) t.imgUnder += 1;
396
+ rec.surf = over(bg, under);
397
+ rec.surfTok = t;
398
+ rec.surfEl = el;
399
+ rec.ownArea = area;
400
+ }
401
+ if (bgImg && bgImg !== 'none') {
402
+ if (bgImg.includes('gradient')) {
403
+ if (ecRow) ecRow[7] = colorsIn(bgImg).map(keyOf).join(',');
404
+ for (const c of colorsIn(bgImg)) { const t = tok(R.gradient, c); flags(t); t.area += area / 2; t.area1 += area1 / 2; }
405
+ const gc = colorsIn(bgImg);
406
+ if (gc.length) {
407
+ let avg = [0, 0, 0, 1];
408
+ for (const c of gc) { const o = over(c, pr.surf); avg[0] += o[0] / gc.length; avg[1] += o[1] / gc.length; avg[2] += o[2] / gc.length; }
409
+ rec.surf = avg; rec.imgBg = false;
410
+ }
411
+ }
412
+ const urls = [...bgImg.matchAll(/url\(\s*["']?([^"')]+)["']?\s*\)/g)].map(m => m[1]);
413
+ if (urls.length && vis && area > 64) {
414
+ images.push({el, kind: 'bg', src: urls[0], x, y, w, h, area, area1, rec});
415
+ if (area > 0.02 * vw * vh) rec.imgBg = true;
416
+ }
417
+ }
418
+ if (ecRow) ecRow[8] = keyOf(rec.surf.map((v, i) => i < 3 ? Math.round(v) : 1));
419
+ // text
420
+ const txt = ownText(el);
421
+ if (txt) {
422
+ const col = parseColor(cs.color);
423
+ if (col) {
424
+ const t = tok(R.text, col);
425
+ if (ecRow) ecRow[3] = keyOf(col);
426
+ flags(t);
427
+ const n = txt.length;
428
+ t.chars += vis ? n : 0;
429
+ const fs = parseFloat(cs.fontSize) || 16;
430
+ t.fsSum += fs * n; t.fsMax = Math.max(t.fsMax, vis ? fs : 0);
431
+ if ((parseInt(cs.fontWeight) || 400) >= 600) t.bold += n;
432
+ t.area += vis ? Math.min(area, n * fs * fs * 0.5) : 0;
433
+ t.area1 += area1 ? Math.min(area1, n * fs * fs * 0.5) : 0;
434
+ const under = rec.surf;
435
+ const ul = oklab(under);
436
+ const wgt = vis ? n : 0.1;
437
+ t.under[0] += ul[0] * wgt; t.under[1] += ul[1] * wgt; t.under[2] += ul[2] * wgt; t.under[3] += wgt;
438
+ const fg = over(col, under);
439
+ const l1 = relLum(fg), l2 = relLum(under);
440
+ t.contrastSum += ((Math.max(l1, l2) + 0.05) / (Math.min(l1, l2) + 0.05)) * wgt; t.contrastW += wgt;
441
+ if (rec.imgBg) t.imgUnder += 1;
442
+ if (vis) addRel('tb', t, rec.surfTok, n);
443
+ // reverse: text sitting on this surface
444
+ const st = rec.surfTok;
445
+ const tl = oklab(col);
446
+ st.over[0] += tl[0] * wgt; st.over[1] += tl[1] * wgt; st.over[2] += tl[2] * wgt; st.over[3] += wgt;
447
+ if (bg) addRel('se', t, tok(R.bg, bg), n);
448
+ }
449
+ }
450
+ // borders
451
+ if (!isSvg) {
452
+ const seen = new Set();
453
+ for (const side of ['Top', 'Right', 'Bottom', 'Left']) {
454
+ const bw = parseFloat(cs['border' + side + 'Width']);
455
+ if (!bw || cs['border' + side + 'Style'] === 'none' || cs['border' + side + 'Style'] === 'hidden') continue;
456
+ const bc = parseColor(cs['border' + side + 'Color']);
457
+ if (!bc) continue;
458
+ const k = keyOf(bc);
459
+ const len = (side === 'Top' || side === 'Bottom') ? w : h;
460
+ if (seen.has(k)) { const t = tok(R.border, bc); t.area += vis ? bw * len : 0; continue; }
461
+ seen.add(k);
462
+ if (ecRow && !ecRow[4]) ecRow[4] = k;
463
+ const t = tok(R.border, bc);
464
+ flags(t);
465
+ t.area += vis ? bw * len : 0;
466
+ const ul = oklab(pr.surf);
467
+ t.under[0] += ul[0]; t.under[1] += ul[1]; t.under[2] += ul[2]; t.under[3] += 1;
468
+ if (vis) addRel('rb', t, pr.surfTok, bw * len);
469
+ if (bg) addRel('se', t, tok(R.bg, bg), 1);
470
+ }
471
+ const ow = parseFloat(cs.outlineWidth);
472
+ if (ow && cs.outlineStyle !== 'none') {
473
+ const oc = parseColor(cs.outlineColor);
474
+ if (oc) { const t = tok(R.border, oc); flags(t); t.area += vis ? ow * 2 * (w + h) : 0; }
475
+ }
476
+ // shadows
477
+ for (const sh of [cs.boxShadow, cs.textShadow]) {
478
+ if (!sh || sh === 'none') continue;
479
+ for (const c of colorsIn(sh)) { const t = tok(R.shadow, c); flags(t); t.area += vis ? Math.min(area, 2 * (w + h) * 4) : 0;
480
+ const ul = oklab(pr.surf); t.under[0] += ul[0]; t.under[1] += ul[1]; t.under[2] += ul[2]; t.under[3] += 1; }
481
+ }
482
+ }
483
+ // svg paint
484
+ if (isSvg) {
485
+ const ltag = tag.toLowerCase();
486
+ if (ltag === 'svg' && !(parent instanceof SVGElement) && vis && area >= 400) {
487
+ const meta = ((el.getAttribute('class') || '') + ' ' + (el.id || '') + ' ' + (el.getAttribute('aria-label') || '') + ' ' + ((el.querySelector('title') || {}).textContent || '')).toLowerCase();
488
+ const useExt = [...el.querySelectorAll('use')].some(u => { const h = u.getAttribute('href') || u.getAttribute('xlink:href') || ''; return h && !h.startsWith('#'); });
489
+ if (rec.inLink || (region & 1) || /logo|brand|wordmark/.test(meta) || useExt || area >= 4000)
490
+ images.push({el, kind: 'svg', src: '', x, y, w, h, area, area1, rec, svgMeta: meta, useExt});
491
+ }
492
+ if (SVG_SHAPES.has(ltag) && ltag !== 'svg' && ltag !== 'g') {
493
+ const fc = cs.fill && cs.fill.startsWith('url') ? null : parseColor(cs.fill);
494
+ const sc = cs.stroke && cs.stroke.startsWith('url') ? null : parseColor(cs.stroke);
495
+ const sw = parseFloat(cs.strokeWidth) || 0;
496
+ const fo = parseFloat(cs.fillOpacity); const so = parseFloat(cs.strokeOpacity);
497
+ if (ecRow && fc) ecRow[5] = keyOf(fc);
498
+ if (ecRow && sc) ecRow[6] = keyOf(sc);
499
+ if (fc && fo !== 0) { const t = tok(R.fill, fc); flags(t); t.area += area * 0.5; t.area1 += area1 * 0.5;
500
+ const ul = oklab(rec.surf); t.under[0] += ul[0]; t.under[1] += ul[1]; t.under[2] += ul[2]; t.under[3] += 1; if (vis) addRel('tb', t, rec.surfTok, Math.sqrt(area + 1)); }
501
+ if (sc && sw > 0 && so !== 0) { const t = tok(R.stroke, sc); flags(t); t.area += vis ? sw * (w + h) : 0;
502
+ const ul = oklab(rec.surf); t.under[0] += ul[0]; t.under[1] += ul[1]; t.under[2] += ul[2]; t.under[3] += 1; if (vis) addRel('tb', t, rec.surfTok, Math.sqrt(area + 1)); }
503
+ }
504
+ }
505
+ // images
506
+ if (tag === 'IMG' && vis && area > 16) {
507
+ images.push({el, kind: 'img', src: el.currentSrc || el.src, x, y, w, h, area, area1, rec});
508
+ } else if ((tag === 'IMAGE' || tag === 'image') && vis) {
509
+ const href = el.getAttribute('href') || el.getAttribute('xlink:href');
510
+ if (href) images.push({el, kind: 'svgimage', src: new URL(href, location.href).href, x, y, w, h, area, area1, rec});
511
+ }
512
+ // pseudo elements
513
+ if (opts.pseudo && vis && !isSvg && els.length < 20000) {
514
+ for (const pseudo of ['::before', '::after']) {
515
+ const ps = getComputedStyle(el, pseudo);
516
+ const ct = ps.content;
517
+ if (!ct || ct === 'none' || ct === 'normal') continue;
518
+ const pb = parseColor(ps.backgroundColor);
519
+ const pw = parseFloat(ps.width) || 0, ph = parseFloat(ps.height) || 0;
520
+ const parea = Math.min(area, pw * ph);
521
+ if (pb) { const t = tok(R.bg, pb); flags(t); t.area += parea; addRel('bb', t, rec.surfTok, parea); }
522
+ if (ct !== '""' && ct !== "''") {
523
+ const pc = parseColor(ps.color);
524
+ if (pc) { const t = tok(R.text, pc); flags(t); t.chars += 1; const ul = oklab(rec.surf); t.under[0] += ul[0]; t.under[1] += ul[1]; t.under[2] += ul[2]; t.under[3] += 1; addRel('tb', t, rec.surfTok, 1); }
525
+ }
526
+ const bw = parseFloat(ps.borderTopWidth) || parseFloat(ps.borderBottomWidth) || parseFloat(ps.borderLeftWidth);
527
+ if (bw) { const bc = parseColor(ps.borderTopColor !== 'rgba(0, 0, 0, 0)' ? ps.borderTopColor : ps.borderBottomColor); if (bc) { const t = tok(R.border, bc); flags(t); t.area += bw * Math.max(pw, ph); } }
528
+ const psh = ps.boxShadow;
529
+ if (psh && psh !== 'none') for (const c of colorsIn(psh)) { const t = tok(R.shadow, c); flags(t); }
530
+ if (ps.backgroundImage && ps.backgroundImage.includes('gradient')) for (const c of colorsIn(ps.backgroundImage)) { const t = tok(R.gradient, c); flags(t); t.area += parea / 2; }
531
+ }
532
+ }
533
+ }
534
+
535
+ // own area for surfaces: subtract direct child surfaces
536
+ // (approximation: walk surfaces; area already includes descendants)
537
+ const childSurfArea = new Map();
538
+ for (const el of els) {
539
+ const r = info.get(el);
540
+ if (!r || !r.surfEl || r.surfEl !== el || !r.ownArea) continue;
541
+ // find nearest ancestor surface element
542
+ let p = parentOf.get(el);
543
+ while (p && !(info.get(p) && info.get(p).surfEl === p)) p = parentOf.get(p);
544
+ const key = p || 'root';
545
+ childSurfArea.set(key, (childSurfArea.get(key) || 0) + r.ownArea);
546
+ }
547
+ const ownAreaTok = new Map();
548
+ let rootOwn = Math.max(0, docW * docH - (childSurfArea.get('root') || 0));
549
+ ownAreaTok.set(rootTok, rootOwn);
550
+ for (const el of els) {
551
+ const r = info.get(el);
552
+ if (!r || r.surfEl !== el || !r.ownArea) continue;
553
+ const own = Math.max(0, r.ownArea - (childSurfArea.get(el) || 0));
554
+ ownAreaTok.set(r.surfTok, (ownAreaTok.get(r.surfTok) || 0) + own);
555
+ }
556
+
557
+ // rule-only colors
558
+ for (const rc of sheetScan.ruleColors) {
559
+ const t = tok(rc.role, rc.c);
560
+ t.rules += rc.n; t.state += rc.state;
561
+ }
562
+ // root token area = canvas own area
563
+ rootTok.area += rootOwn;
564
+
565
+ // ---------- darkness measurement ----------
566
+ let bgLsum = 0, bgW = 0, darkW = 0;
567
+ for (const [t, own] of ownAreaTok) {
568
+ if (t.role !== R.bg) continue;
569
+ const eff = over(t.c, t.root ? canvasColor : (t.under[3] ? labToApproxRgb(t.under) : canvasColor));
570
+ const L = oklab(eff)[0];
571
+ bgLsum += L * own; bgW += own; if (L < 0.45) darkW += own;
572
+ }
573
+ let txLsum = 0, txW = 0;
574
+ for (const t of tokens.values()) if (t.role === R.text && t.chars) { txLsum += oklab(t.c)[0] * t.chars; txW += t.chars; }
575
+ const darkness = {bgL: bgW ? bgLsum / bgW : 1, darkFrac: bgW ? darkW / bgW : 0, textL: txW ? txLsum / txW : 0, rootL: oklab(rootSurface)[0]};
576
+
577
+ // ---------- token selection & features ----------
578
+ let list = [...tokens.values()];
579
+ const score = (t) => (t.root ? 1e12 : 0) + Math.log1p(t.area) * 3 + Math.log1p(t.chars) * 2 + Math.log1p(t.nVis) * 2 + Math.log1p(t.n) + Math.log1p(t.rules) * 0.7 + Math.log1p(t.state);
580
+ list.sort((a, b) => score(b) - score(a));
581
+ const dropped = Math.max(0, list.length - opts.maxTokens);
582
+ list = list.slice(0, opts.maxTokens);
583
+ list.forEach((t, i) => t.i = i);
584
+
585
+ // bg lightness rank among visible bg tokens
586
+ const bgVis = list.filter(t => t.role === R.bg && t.area > 0).map(t => oklab(t.c)[0]).sort((a, b) => a - b);
587
+ const rankL = (L) => { if (!bgVis.length) return 0.5; let i = 0; while (i < bgVis.length && bgVis[i] < L) i++; return i / bgVis.length; };
588
+ const totalChars = list.reduce((s, t) => s + (t.role === R.text ? t.chars : 0), 0) || 1;
589
+
590
+ const out = [];
591
+ for (const t of list) {
592
+ const lab = oklab(t.c);
593
+ const C = Math.hypot(lab[1], lab[2]);
594
+ const under = t.under[3] ? [t.under[0] / t.under[3], t.under[1] / t.under[3], t.under[2] / t.under[3]] : oklab(canvasColor);
595
+ const overL = t.over[3] ? [t.over[0] / t.over[3], t.over[1] / t.over[3], t.over[2] / t.over[3]] : [0, 0, 0];
596
+ const n = Math.max(1, t.n);
597
+ const own = ownAreaTok.get(t) || 0;
598
+ const f = [
599
+ lab[0], lab[1] * 4, lab[2] * 4, C * 4, t.c[3], C > 0.02 ? lab[1] / C : 0, C > 0.02 ? lab[2] / C : 0,
600
+ t.c[0] / 255, t.c[1] / 255, t.c[2] / 255,
601
+ Math.log1p(t.area / 1000) / 10, Math.sqrt(Math.min(1, t.area / pageArea)), Math.sqrt(Math.min(1, t.area1 / (vw * vh))), Math.sqrt(Math.min(1, own / pageArea)),
602
+ Math.log1p(t.n) / 8, Math.log1p(t.nVis) / 8, Math.log1p(t.rules) / 6, t.rules ? t.state / (t.rules) : 0, (t.n === 0 && t.rules > 0) ? 1 : 0,
603
+ Math.log1p(t.chars) / 10, t.chars / totalChars, t.chars ? (t.fsSum / t.chars) / 32 : 0, t.fsMax / 48, t.chars ? t.bold / t.chars : 0,
604
+ t.root, t.inter / n, t.link / n, t.head / n, t.code / n, t.form / n, t.header / n, t.nav / n, t.footer / n, t.fixed / n,
605
+ under[0], under[1] * 4, under[2] * 4, t.contrastW ? Math.log(t.contrastSum / t.contrastW) / 3 : 0, Math.min(1, t.imgUnder / n),
606
+ overL[0], overL[1] * 4, overL[2] * 4, t.over[3] ? 1 : 0,
607
+ t.role === R.bg ? rankL(lab[0]) : 0, t.nVis ? Math.min(1, t.depthSum / t.nVis / 30) : 0,
608
+ t.ws && t.area ? Math.min(1, t.ws / t.area / docW) : 0, t.hs && t.area ? Math.min(1, t.hs / t.area / vh) : 0,
609
+ t.ys && t.area ? Math.min(1, t.ys / t.area / docH) : 0,
610
+ t.media / n,
611
+ ];
612
+ for (let q = 0; q < f.length; q++) if (!Number.isFinite(f[q])) f[q] = 0; // JSON turns NaN into null; never feed it to the model
613
+ const o = {k: t.k, role: t.role, rgba: [t.c[0], t.c[1], t.c[2], +t.c[3].toFixed(3)], f};
614
+ if (opts.describe) {
615
+ o.d = {area: Math.round(t.area), area1: Math.round(t.area1), own: Math.round(own), n: t.n, nVis: t.nVis, rules: t.rules, state: t.state, chars: t.chars,
616
+ fs: t.chars ? +(t.fsSum / t.chars).toFixed(1) : 0, root: t.root, inter: t.inter, link: t.link, head: t.head, header: t.header, nav: t.nav, footer: t.footer,
617
+ contrast: t.contrastW ? +(t.contrastSum / t.contrastW).toFixed(2) : null, media: t.media, ex: t.ex};
618
+ }
619
+ out.push(o);
620
+ }
621
+ const relOut = [];
622
+ for (const [k, w] of rel) {
623
+ const [type, a, b] = k.split('|');
624
+ const ta = tokens.get(a), tb = tokens.get(b);
625
+ if (!ta || !tb || ta.i < 0 || tb.i < 0) continue;
626
+ relOut.push([['tb', 'bb', 'rb', 'se'].indexOf(type), ta.i, tb.i, +Math.log1p(w).toFixed(3)]);
627
+ }
628
+
629
+ // ---------- images ----------
630
+ const imgSeen = new Map();
631
+ let svgCount = 0;
632
+ // logo-like images first, then by area
633
+ const logoish = (im) => (im.rec.link && im.rec.inLink ? 1 : 0) + ((im.rec.region & 1) ? 1 : 0) + (im.kind === 'svg' && /logo|brand/.test(im.svgMeta || '') ? 2 : 0);
634
+ images.sort((a, b) => (logoish(b) - logoish(a)) || (b.area - a.area));
635
+ const imgOut = [];
636
+ for (const im of images) {
637
+ if (imgOut.length >= opts.maxImages) break;
638
+ if (im.kind !== 'svg' && (!im.src || im.src.startsWith('data:image/gif;base64,R0lGOD'))) continue;
639
+ if (im.kind === 'svg') { im.src = 'svg:' + svgCount++; if (opts.markSvg) im.el.setAttribute('data-nadir-svg', im.src); }
640
+ const key = im.src;
641
+ if (imgSeen.has(key)) { const o = imgOut[imgSeen.get(key)]; o.count++; o.areaSum += im.area; continue; }
642
+ const el = im.el;
643
+ const a = im.rec.link;
644
+ const href = a ? (a.getAttribute('href') || '') : '';
645
+ let homeLink = 0;
646
+ try { if (a) { const u = new URL(href, location.href); homeLink = (u.origin === location.origin && (u.pathname === '/' || u.pathname === '')) ? 1 : 0; } } catch (e) {}
647
+ 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();
648
+ const ext = im.kind === 'svg' ? 'svg' : (im.src.split('?')[0].match(/\.(svg|png|jpe?g|webp|gif|avif|ico)$/i) || [, ''])[1].toLowerCase();
649
+ const under = oklab(im.rec.surf);
650
+ 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,
651
+ 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,
652
+ 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,
653
+ ext, underTok: im.rec.surfTok && im.rec.surfTok.i >= 0 ? im.rec.surfTok.i : -1, under,
654
+ 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,
655
+ alt: opts.describe ? ((el.getAttribute && el.getAttribute('alt')) || '').slice(0, 60) : undefined,
656
+ stats: null, thumb: null};
657
+ if (im.kind === 'img' && el.complete === false) o.loading = 1;
658
+ imgSeen.set(key, imgOut.length);
659
+ imgOut.push(o);
660
+ }
661
+ if (opts.thumbs) await imageStats(imgOut, opts);
662
+ for (const o of imgOut) o.f = imageFeatures(o, vw, vh, docH);
663
+ const res = {
664
+ v: 5, url: location.href, title: document.title.slice(0, 200), vw, vh, docH, docW,
665
+ scheme: schemeInfo(sheetScan), darkness, canvasDark,
666
+ tokens: out, rel: relOut, images: imgOut,
667
+ ec: ec || undefined,
668
+ stats: {ads: adsSkipped, elements: els.length, visible: visibleCount, tokensTotal: tokens.size, tokensDropped: dropped, ms: Math.round(performance.now() - t0)},
669
+ };
670
+ return res;
671
+ }
672
+
673
+ function labToApproxRgb(u) {
674
+ // inverse OKLab for averaged under-colour
675
+ const L = u[0] / u[3], a = u[1] / u[3], b = u[2] / u[3];
676
+ const l_ = L + 0.3963377774 * a + 0.2158037573 * b;
677
+ const m_ = L - 0.1055613458 * a - 0.0638541728 * b;
678
+ const s_ = L - 0.0894841775 * a - 1.2914855480 * b;
679
+ const l = l_ ** 3, m = m_ ** 3, s = s_ ** 3;
680
+ const r = 4.0767416621 * l - 3.3077115913 * m + 0.2309699292 * s;
681
+ const g = -1.2684380046 * l + 2.6097574011 * m - 0.3413193965 * s;
682
+ const bb = -0.0041960863 * l - 0.7034186147 * m + 1.7076147010 * s;
683
+ const enc = (x) => { x = Math.max(0, Math.min(1, x)); return 255 * (x <= 0.0031308 ? 12.92 * x : 1.055 * Math.pow(x, 1 / 2.4) - 0.055); };
684
+ return [enc(r), enc(g), enc(bb), 1];
685
+ }
686
+
687
+ // ---------- image pixels ----------
688
+ function blobToDataURL(blob) {
689
+ return new Promise((res) => { const fr = new FileReader(); fr.onload = () => res(fr.result); fr.onerror = () => res(null); fr.readAsDataURL(blob); });
690
+ }
691
+ function loadViaImg(url) {
692
+ return new Promise((res) => {
693
+ const img = new Image(); img.decoding = 'async';
694
+ const done = (ok) => res(ok && img.naturalWidth ? img : null);
695
+ img.onload = () => done(true); img.onerror = () => done(false); setTimeout(() => done(false), 4000);
696
+ img.src = url;
697
+ });
698
+ }
699
+ // Decode without loading a URL into the page where possible (page CSP often blocks blob: images).
700
+ async function loadImage(src, fetchImage) {
701
+ if (fetchImage && !src.startsWith('data:')) {
702
+ const blob = await fetchImage(src);
703
+ if (!blob || !blob.size) return null;
704
+ const isSvg = /svg/.test(blob.type) || /\.svg(\?|#|$)/i.test(src);
705
+ if (!isSvg) { try { return await createImageBitmap(blob); } catch (e) {} }
706
+ const du = await blobToDataURL(isSvg ? new Blob([blob], {type: 'image/svg+xml'}) : blob);
707
+ return du ? await loadViaImg(du) : null;
708
+ }
709
+ return await loadViaImg(src);
710
+ }
711
+
712
+ async function imageStats(list, opts) {
713
+ const jobs = list.map(async (o) => {
714
+ try {
715
+ const img = await loadImage(o.src, opts.fetchImage);
716
+ if (!img) return;
717
+ const IW = img.naturalWidth || img.width, IH = img.naturalHeight || img.height;
718
+ o.nw = o.nw || IW; o.nh = o.nh || IH;
719
+ statsFromSource(o, img, IW, IH, opts);
720
+ if (img.close) img.close();
721
+ } catch (e) { o.err = String(e).slice(0, 80); }
722
+ });
723
+ await Promise.all(jobs);
724
+ }
725
+ function drawThumbSrc(src, iw, ih, size) {
726
+ const c = document.createElement('canvas'); c.width = size; c.height = size;
727
+ const ctx = c.getContext('2d', {willReadFrequently: true});
728
+ const s = Math.min(size / iw, size / ih);
729
+ const dw = Math.max(1, Math.round(iw * s)), dh = Math.max(1, Math.round(ih * s));
730
+ ctx.drawImage(src, Math.floor((size - dw) / 2), Math.floor((size - dh) / 2), dw, dh);
731
+ return {c, ctx, dw, dh};
732
+ }
733
+ function statsFromSource(o, img, IW, IH, opts) {
734
+ const S = opts.thumbSize;
735
+ const {ctx, dw, dh} = drawThumbSrc(img, IW, IH, S);
736
+ const d = ctx.getImageData(0, 0, S, S).data;
737
+ o.thumb = Array.from(d); // RGBA bytes, letterboxed (transparent pad)
738
+ let n = 0, tr = 0, semi = 0, Ls = 0, L2 = 0, Cs = 0, white = 0, black = 0, ringA = 0, ringL = 0, ringN = 0, total = 0;
739
+ const x0 = Math.floor((S - dw) / 2), y0 = Math.floor((S - dh) / 2);
740
+ const qset = new Set();
741
+ for (let y = y0; y < y0 + dh; y++) for (let x = x0; x < x0 + dw; x++) {
742
+ const i = (y * S + x) * 4; const a = d[i + 3] / 255; total++;
743
+ const ring = (x === x0 || y === y0 || x === x0 + dw - 1 || y === y0 + dh - 1);
744
+ if (ring) { ringA += a; ringN++; }
745
+ if (a < 0.1) { tr++; continue; }
746
+ if (a < 0.9) semi++;
747
+ const lab = oklab([d[i], d[i + 1], d[i + 2]]);
748
+ const C = Math.hypot(lab[1], lab[2]);
749
+ n++; Ls += lab[0]; L2 += lab[0] * lab[0]; Cs += C;
750
+ if (lab[0] > 0.93 && C < 0.03) white++;
751
+ if (lab[0] < 0.3 && C < 0.04) black++;
752
+ if (ring) ringL += lab[0];
753
+ qset.add((d[i] >> 5) << 6 | (d[i + 1] >> 5) << 3 | (d[i + 2] >> 5));
754
+ }
755
+ const mL = n ? Ls / n : 0;
756
+ o.stats = {transp: total ? tr / total : 0, semi: total ? semi / total : 0, meanL: mL, stdL: n ? Math.sqrt(Math.max(0, L2 / n - mL * mL)) : 0,
757
+ meanC: n ? Cs / n : 0, white: n ? white / n : 0, black: n ? black / n : 0, ringA: ringN ? ringA / ringN : 0,
758
+ ringL: ringN ? ringL / Math.max(1, ringN * (ringA / ringN || 1)) : 0, colors: qset.size / 512};
759
+ if (opts.bigThumbSize) o.big = drawThumbSrc(img, IW, IH, opts.bigThumbSize).c.toDataURL('image/png');
760
+ }
761
+ function imageFeatures(o, vw, vh, docH) {
762
+ const s = o.stats || {};
763
+ const ext = o.ext;
764
+ return sanitize([
765
+ Math.log1p(o.w) / 8, Math.log1p(o.h) / 8, Math.log1p(o.areaSum / 100) / 12, Math.min(1, (o.w * o.h) / (vw * vh)), Math.log(Math.max(1, o.w) / Math.max(1, o.h)) / 3,
766
+ Math.log1p(o.nw) / 9, Math.log1p(o.nh) / 9, Math.min(1, o.y / Math.max(1, docH)), o.y < vh ? 1 : 0, Math.log1p(o.count) / 3,
767
+ o.homeLink, o.inLink, o.header, o.footer, o.logo, o.icon, o.avatar,
768
+ ext === 'svg' ? 1 : 0, ext === 'png' ? 1 : 0, (ext === 'jpg' || ext === 'jpeg') ? 1 : 0, ext === 'webp' ? 1 : 0, ext === 'gif' ? 1 : 0, o.kind === 'bg' ? 1 : 0,
769
+ o.under ? o.under[0] : 1, o.under ? o.under[1] * 4 : 0, o.under ? o.under[2] * 4 : 0,
770
+ 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,
771
+ o.kind === 'svg' ? 1 : 0, o.useExt || 0,
772
+ ]);
773
+ }
774
+ const sanitize = (f) => f.map((v) => (Number.isFinite(+v) ? +v : 0));
775
+
776
+ // Compute pixel stats for a list of image entries outside the page (helper page / extension offscreen document).
777
+ // opts.fetchImage(src) -> Blob; opts.shot: dataURL of the viewport screenshot (for inline-SVG crops).
778
+ async function defaultFetch(src) {
779
+ const ac = new AbortController(); const t = setTimeout(() => ac.abort(), 6000);
780
+ try { const r = await fetch(src, {signal: ac.signal, credentials: 'omit'}); if (!r.ok) return null; const b = await r.blob(); return b.size > 8e6 ? null : b; }
781
+ catch (e) { return null; } finally { clearTimeout(t); }
782
+ }
783
+ async function imageStatsStandalone(list, opts) {
784
+ opts = Object.assign({thumbSize: 32, bigThumbSize: 0, concurrency: 8, fetchImage: defaultFetch}, opts || {});
785
+ let shotImg = null;
786
+ if (opts.shot && list.some(o => o.kind === 'svg')) shotImg = await loadViaImg(opts.shot);
787
+ // inline-SVG crops first (local, fast), then fetched images with limited concurrency
788
+ for (const o of list) {
789
+ if (o.kind !== 'svg' || !shotImg) continue;
790
+ const dpr = opts.dpr || 1;
791
+ if (o.vy < 0 || (o.vy + o.h) * dpr > shotImg.naturalHeight + 1 || o.w < 2 || o.h < 2) continue;
792
+ cropStats(o, shotImg, opts);
793
+ }
794
+ const raster = list.filter(o => o.kind !== 'svg');
795
+ let k = 0;
796
+ const run = async () => { while (k < raster.length) { const o = raster[k++]; await imageStats([o], opts); } };
797
+ await Promise.all(Array.from({length: opts.concurrency}, run));
798
+ return list.map(o => ({stats: o.stats || null, thumb: o.thumb || null, big: o.big || null, nw: o.nw, nh: o.nh, err: o.err}));
799
+ }
800
+ // inline SVGs: crop the rendered element from a screenshot; alpha is keyed off the surface colour beneath it
801
+ function cropStats(o, shot, opts) {
802
+ const S = opts.thumbSize;
803
+ const c = document.createElement('canvas'); c.width = Math.max(1, Math.round(o.w)); c.height = Math.max(1, Math.round(o.h));
804
+ const cx = c.getContext('2d', {willReadFrequently: true});
805
+ const dpr = opts.dpr || 1;
806
+ cx.drawImage(shot, o.vx * dpr, o.vy * dpr, o.w * dpr, o.h * dpr, 0, 0, c.width, c.height);
807
+ const src = cx.getImageData(0, 0, c.width, c.height);
808
+ const sl = oklab(o.surf || [255, 255, 255]);
809
+ const d = src.data;
810
+ for (let i = 0; i < d.length; i += 4) {
811
+ const lab = oklab([d[i], d[i + 1], d[i + 2]]);
812
+ const dist = Math.hypot(lab[0] - sl[0], lab[1] - sl[1], lab[2] - sl[2]);
813
+ d[i + 3] = Math.round(255 * Math.max(0, Math.min(1, (dist - 0.02) / 0.1)));
814
+ }
815
+ cx.putImageData(src, 0, 0);
816
+ o.nw = c.width; o.nh = c.height;
817
+ statsFromSource(o, c, c.width, c.height, opts);
818
+ }
819
+
820
+ globalThis.__nadirImageStats = imageStatsStandalone;
821
+ globalThis.__nadirImageFeatures = imageFeatures;
822
+ globalThis.__nadirExtract = extract;
823
+ globalThis.__nadirOklab = oklab;
824
+ globalThis.__nadirParseColor = parseColor;
825
+ })();
runtime/nadir-infer.js ADDED
@@ -0,0 +1,193 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ // Nadir inference glue: page description (from extract.js + image stats) -> model inputs -> dark mapping.
2
+ // Mirrors train.py collate() for a single page. Works in browsers (offscreen document) and Node (parity tests).
3
+ (function (root) {
4
+ 'use strict';
5
+ const MAX_TOK = 256, MAX_IMG = 32, FT = 49, FI = 39, K = 11;
6
+ const TREAT = ['keep', 'dim', 'invert', 'light', 'plate'];
7
+ const FG_ROLES = new Set([1, 5, 6]); // text, svg fill, svg stroke
8
+ const lin = (u) => { u /= 255; return u <= 0.04045 ? u / 12.92 : Math.pow((u + 0.055) / 1.055, 2.4); };
9
+ function oklab(r, g, b) {
10
+ r = lin(r); g = lin(g); b = lin(b);
11
+ const l = Math.cbrt(0.4122214708 * r + 0.5363325363 * g + 0.0514459929 * b);
12
+ const m = Math.cbrt(0.2119034982 * r + 0.6806995451 * g + 0.1073969566 * b);
13
+ const s = Math.cbrt(0.0883024619 * r + 0.2817188376 * g + 0.6299787005 * b);
14
+ return [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];
15
+ }
16
+ function labToRgb(L, a, b) {
17
+ const l_ = L + 0.3963377774 * a + 0.2158037573 * b, m_ = L - 0.1055613458 * a - 0.0638541728 * b, s_ = L - 0.0894841775 * a - 1.2914855480 * b;
18
+ const l = l_ ** 3, m = m_ ** 3, s = s_ ** 3;
19
+ const enc = (x) => { x = Math.max(0, Math.min(1, x)); return Math.round(255 * (x <= 0.0031308 ? 12.92 * x : 1.055 * Math.pow(x, 1 / 2.4) - 0.055)); };
20
+ return [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)];
21
+ }
22
+ const relLum = (rgb) => 0.2126 * lin(rgb[0]) + 0.7152 * lin(rgb[1]) + 0.0722 * lin(rgb[2]);
23
+ const contrast = (a, b) => { const x = relLum(a), y = relLum(b); return (Math.max(x, y) + 0.05) / (Math.min(x, y) + 0.05); };
24
+ 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') : '');
25
+
26
+ function buildInputs(page) {
27
+ const toks = page.tokens.slice(0, MAX_TOK);
28
+ const N = Math.max(1, toks.length);
29
+ const imgs = page.images.slice(0, MAX_IMG);
30
+ const Mreal = imgs.length, M = Math.max(1, Mreal);
31
+ const T = 1 + N + M;
32
+ const tok_f = new Float32Array(N * FT), role = new BigInt64Array(N), lab = new Float32Array(N * 3);
33
+ toks.forEach((t, i) => {
34
+ tok_f.set(t.f.slice(0, FT).map((v) => (Number.isFinite(v) ? v : 0)), i * FT); role[i] = BigInt(t.role);
35
+ const o = oklab(t.rgba[0], t.rgba[1], t.rgba[2]); lab.set(o, i * 3);
36
+ });
37
+ const img_f = new Float32Array(M * FI), thumbs = new Float32Array(M * 4 * 32 * 32);
38
+ imgs.forEach((im, j) => {
39
+ img_f.set(im.f.slice(0, FI).map((v) => (Number.isFinite(v) ? v : 0)), j * FI);
40
+ if (im.thumb && im.thumb.length === 4096) {
41
+ for (let y = 0; y < 32; y++) for (let x = 0; x < 32; x++) for (let c = 0; c < 4; c++)
42
+ thumbs[j * 4096 + c * 1024 + y * 32 + x] = im.thumb[(y * 32 + x) * 4 + c] / 255;
43
+ }
44
+ });
45
+ const st = page.stats || {};
46
+ const page_f = new Float32Array([page.vw / 1280, Math.min(8, page.docH / page.vh) / 8, toks.length / MAX_TOK, Mreal / MAX_IMG,
47
+ page.darkness.bgL, page.darkness.textL, page.canvasDark ? 1 : 0, Math.min(1, (st.tokensDropped || 0) / 200)]);
48
+ // relations
49
+ const rel = new Float32Array(T * T * K);
50
+ const n = toks.length;
51
+ const R = new Float32Array(4 * (n + 1) * (n + 1));
52
+ const P = n + 1;
53
+ const tbBest = new Map();
54
+ for (const r of page.rel) {
55
+ const [typ, a, b, w] = r;
56
+ if (a >= n || b >= n) continue;
57
+ const idx = typ * P * P + (a + 1) * P + (b + 1);
58
+ if (w > R[idx]) R[idx] = w;
59
+ if (typ === 0) { const cur = tbBest.get(a); if (!cur || w > cur[0]) tbBest.set(a, [w, b]); }
60
+ }
61
+ for (let k = 0; k < 4; k++) for (let a = 0; a < P; a++) {
62
+ let mx = 0; for (let b = 0; b < P; b++) mx = Math.max(mx, R[k * P * P + a * P + b]);
63
+ const den = mx + 1e-6;
64
+ for (let b = 0; b < P; b++) {
65
+ const v = R[k * P * P + a * P + b] / den;
66
+ if (!v) continue;
67
+ rel[(a * T + b) * K + k] = v; // a -> b
68
+ rel[(b * T + a) * K + 4 + k] = v; // transpose channel
69
+ }
70
+ }
71
+ imgs.forEach((im, j) => {
72
+ const u = im.underTok;
73
+ if (u != null && u >= 0 && u < n) { rel[((1 + N + j) * T + (1 + u)) * K + 8] = 1; rel[((1 + u) * T + (1 + N + j)) * K + 9] = 1; }
74
+ });
75
+ const keyOf = (t) => t.rgba[0] * 65536 + t.rgba[1] * 256 + t.rgba[2];
76
+ const byKey = new Map();
77
+ toks.forEach((t, i) => { const k = keyOf(t); if (!byKey.has(k)) byKey.set(k, []); byKey.get(k).push(i); });
78
+ for (const list of byKey.values()) if (list.length > 1) for (const i of list) for (const j of list) if (i !== j) rel[((1 + i) * T + (1 + j)) * K + 10] = 1;
79
+ const mask = new Uint8Array(T);
80
+ for (let i = 0; i <= n; i++) mask[i] = 1;
81
+ for (let j = 0; j < Mreal; j++) mask[1 + N + j] = 1;
82
+ return {N, M, T, Mreal, tbBest, feeds: {
83
+ tok_f: [tok_f, [1, N, FT]], tok_role: [role, [1, N]], tok_lab: [lab, [1, N, 3]], img_f: [img_f, [1, M, FI]], img_thumb: [thumbs, [1, M, 4, 32, 32]],
84
+ page_f: [page_f, [1, 8]], src: [new BigInt64Array([0n]), [1]], rel: [rel, [1, T, T, K]], mask: [mask, [1, T]],
85
+ }};
86
+ }
87
+
88
+ // outputs: lab Float32Array(N*3), alpha_logit Float32Array(N), img_logits Float32Array(M*5)
89
+ // brand colours keep their hue: a coloured input whose output drifts more than 45 degrees is turned back to its
90
+ // own hue (keeping the predicted lightness and chroma); neutral inputs get at most a faint tint
91
+ function hueGuard(t, o) {
92
+ const i = oklab(t.rgba[0], t.rgba[1], t.rgba[2]), q = oklab(...o.rgb);
93
+ const cin = Math.hypot(i[1], i[2]), cout = Math.hypot(q[1], q[2]);
94
+ if (cin > 0.03 && cout > 0.03) {
95
+ const d = Math.abs(((Math.atan2(q[2], q[1]) - Math.atan2(i[2], i[1]) + Math.PI) % (2 * Math.PI) + 2 * Math.PI) % (2 * Math.PI) - Math.PI);
96
+ if (d <= Math.PI / 4) return;
97
+ const h = Math.atan2(i[2], i[1]);
98
+ o.rgb = labToRgb(q[0], cout * Math.cos(h), cout * Math.sin(h));
99
+ } else if (cin <= 0.03 && cout > 0.06) {
100
+ o.rgb = labToRgb(q[0], q[1] * 0.06 / cout, q[2] * 0.06 / cout);
101
+ }
102
+ }
103
+
104
+ function decode(page, built, out, opts) {
105
+ opts = opts || {};
106
+ const toks = page.tokens.slice(0, MAX_TOK);
107
+ const n = toks.length;
108
+ const outputs = new Array(n);
109
+ const keep = new Set();
110
+ toks.forEach((t, i) => {
111
+ const L = out.lab[i * 3], a = out.lab[i * 3 + 1], b = out.lab[i * 3 + 2];
112
+ const ain = Math.min(1 - 1e-3, Math.max(1e-3, t.rgba[3]));
113
+ const alpha = 1 / (1 + Math.exp(-(out.alpha_logit[i] + Math.log(ain / (1 - ain)))));
114
+ outputs[i] = {rgb: labToRgb(L, a, b), a: t.rgba[3] >= 0.999 && alpha > 0.97 ? 1 : alpha, L};
115
+ // colours over video/photos that already read as dark mode (dark scrims, saturated player accents, light
116
+ // captions and icons) are shown as-is; anything else over media is redesigned like any other colour
117
+ const f = t.f;
118
+ const bgLike = t.role === 0 || t.role === 3 || t.role === 4;
119
+ if (f[48] >= 0.5 && (bgLike ? (f[24] <= 0.5 && f[11] < 0.25 && (f[0] < 0.5 || f[3] > 0.4)) : f[0] > 0.6)) {
120
+ keep.add(i);
121
+ const rgb = t.rgba.slice(0, 3);
122
+ outputs[i] = {rgb, a: t.rgba[3], L: oklab(...rgb)[0]};
123
+ } else hueGuard(t, outputs[i]);
124
+ });
125
+ // safety: page canvas must be dark
126
+ toks.forEach((t, i) => {
127
+ 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; }
128
+ });
129
+ // things that were invisible against their surface stay invisible: gradient stops with exactly a surface's
130
+ // colour (text fades) take that surface's dark colour, and so do borders that matched their surface
131
+ const keyOf = (t) => t.rgba[0] * 65536 + t.rgba[1] * 256 + t.rgba[2];
132
+ const surfOf = new Map();
133
+ toks.forEach((t, i) => { if (t.role === 0 && t.rgba[3] >= 0.999 && !keep.has(i) && outputs[i].a >= 0.999 && !surfOf.has(keyOf(t))) surfOf.set(keyOf(t), i); });
134
+ toks.forEach((t, i) => {
135
+ if (t.role === 4 && !keep.has(i) && surfOf.has(keyOf(t))) { const j = surfOf.get(keyOf(t)); outputs[i].rgb = outputs[j].rgb.slice(); outputs[i].L = outputs[j].L; }
136
+ });
137
+ const borderOn = new Map();
138
+ for (const [typ, a, b, w] of page.rel) {
139
+ if (typ !== 2 || a >= n || b >= n || toks[a].role !== 2 || toks[b].role !== 0) continue;
140
+ const cur = borderOn.get(a); if (!cur || w > cur[0]) borderOn.set(a, [w, b]);
141
+ }
142
+ for (const [i, [, j]] of borderOn) {
143
+ if (keep.has(i) || keep.has(j)) continue;
144
+ if (contrast(toks[i].rgba, toks[j].rgba) < 1.05 && toks[i].rgba[3] >= 0.5) { outputs[i].rgb = outputs[j].rgb.slice(); outputs[i].L = outputs[j].L; }
145
+ }
146
+ // safety: every foreground colour (text, svg fill/stroke) stays readable on all the surfaces it sits on,
147
+ // never asking for more contrast than the light design had; one lightness per shared token, chosen on a grid
148
+ const surf = new Map();
149
+ for (const [typ, a, b, w] of page.rel) {
150
+ if (typ !== 0 || a >= n || b >= n || toks[b].role !== 0 || outputs[b].a < 0.999) continue;
151
+ if (!surf.has(a)) surf.set(a, []);
152
+ surf.get(a).push([b, w + 0.1]);
153
+ }
154
+ for (const [i, lst] of surf) {
155
+ const t = toks[i];
156
+ if (keep.has(i) || !FG_ROLES.has(t.role) || outputs[i].a < 0.999) continue;
157
+ let need0 = 3;
158
+ 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; }
159
+ if (opts.minContrast) need0 = Math.min(need0, opts.minContrast);
160
+ const reqs = [];
161
+ for (const [j, w] of lst) {
162
+ const need = Math.min(need0, contrast(t.rgba, toks[j].rgba));
163
+ if (need >= 1.5) reqs.push([outputs[j].rgb, need, w]);
164
+ }
165
+ if (!reqs.length) continue;
166
+ const deficit = (rgb) => { let d = 0; for (const [bg, need, w] of reqs) d += w * Math.max(0, Math.log(need) - Math.log(contrast(rgb, bg))); return d; };
167
+ if (deficit(outputs[i].rgb) <= 0) continue;
168
+ const o = oklab(...outputs[i].rgb);
169
+ let W = 0; for (const r of reqs) W += r[2];
170
+ let best = null;
171
+ for (let k = 0; k <= 50; k++) {
172
+ const Lx = k / 50, rgb = labToRgb(Lx, o[1], o[2]);
173
+ const c = deficit(rgb) / W + 0.05 * Math.abs(Lx - o[0]);
174
+ if (!best || c < best[0]) best = [c, rgb, Lx];
175
+ }
176
+ outputs[i].rgb = best[1]; outputs[i].L = best[2];
177
+ }
178
+ const colors = toks.map((t, i) => keep.has(i) ? null : hex(outputs[i].rgb, outputs[i].a));
179
+ const imgs = page.images.slice(0, MAX_IMG);
180
+ const treatments = imgs.map((im, j) => {
181
+ let best = 0;
182
+ for (let c = 1; c < 5; c++) if (out.img_logits[j * 5 + c] > out.img_logits[j * 5 + best]) best = c;
183
+ const t = TREAT[best];
184
+ const u = im.underTok;
185
+ return root.NadirSafeTreatment ? root.NadirSafeTreatment(t, im, u != null && u >= 0 && u < n ? outputs[u].L : null) : t;
186
+ });
187
+ return {colors, keep: [...keep], treatments};
188
+ }
189
+
190
+ const api = {buildInputs, decode, oklab, labToRgb, TREAT, MAX_TOK, MAX_IMG};
191
+ root.NadirInfer = api;
192
+ if (typeof module !== 'undefined') module.exports = api;
193
+ })(typeof globalThis !== 'undefined' ? globalThis : this);
runtime/nadir-map.js ADDED
@@ -0,0 +1,244 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ // Nadir colour mapping runtime: turns per-token dark colours (from the model or a teacher) into a lookup
2
+ // for every colour the page might use. Shared by the extension and the evaluation renderer.
3
+ (function () {
4
+ 'use strict';
5
+ const ROLES = ['bg', 'text', 'border', 'shadow', 'gradient', 'fill', 'stroke'];
6
+ // which token roles answer a query role, in order of preference
7
+ const CHAIN = {
8
+ bg: ['bg', 'gradient', 'shadow'],
9
+ gradient: ['gradient', 'bg', 'shadow'],
10
+ shadow: ['shadow', 'bg', 'gradient'],
11
+ text: ['text', 'fill', 'stroke'],
12
+ fill: ['fill', 'text', 'stroke'],
13
+ stroke: ['stroke', 'fill', 'text'],
14
+ border: ['border', 'bg', 'stroke'],
15
+ };
16
+ const lin = (u) => { u /= 255; return u <= 0.04045 ? u / 12.92 : Math.pow((u + 0.055) / 1.055, 2.4); };
17
+ function oklab(r, g, b) {
18
+ r = lin(r); g = lin(g); b = lin(b);
19
+ const l = Math.cbrt(0.4122214708 * r + 0.5363325363 * g + 0.0514459929 * b);
20
+ const m = Math.cbrt(0.2119034982 * r + 0.6806995451 * g + 0.1073969566 * b);
21
+ const s = Math.cbrt(0.0883024619 * r + 0.2817188376 * g + 0.6299787005 * b);
22
+ return [0.2104542553 * l + 0.7936177850 * m - 0.0040720468 * s,
23
+ 1.9779984951 * l - 2.4285922050 * m + 0.4505937099 * s,
24
+ 0.0259040371 * l + 0.7827717662 * m - 0.8086757660 * s];
25
+ }
26
+ function toRgb(L, a, b) {
27
+ const l_ = L + 0.3963377774 * a + 0.2158037573 * b;
28
+ const m_ = L - 0.1055613458 * a - 0.0638541728 * b;
29
+ const s_ = L - 0.0894841775 * a - 1.2914855480 * b;
30
+ const l = l_ ** 3, m = m_ ** 3, s = s_ ** 3;
31
+ const enc = (x) => { x = Math.max(0, Math.min(1, x)); return Math.round(255 * (x <= 0.0031308 ? 12.92 * x : 1.055 * Math.pow(x, 1 / 2.4) - 0.055)); };
32
+ return [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)];
33
+ }
34
+ const hex2 = (n) => n.toString(16).padStart(2, '0');
35
+ const keyOf = (r, g, b, a) => hex2(r) + hex2(g) + hex2(b) + hex2(Math.round(a * 255));
36
+ function css(rgb, a) {
37
+ return a >= 0.999 ? `#${hex2(rgb[0])}${hex2(rgb[1])}${hex2(rgb[2])}` : `rgba(${rgb[0]}, ${rgb[1]}, ${rgb[2]}, ${+a.toFixed(3)})`;
38
+ }
39
+ function parseHex(h) {
40
+ h = h.replace('#', '');
41
+ return [parseInt(h.slice(0, 2), 16), parseInt(h.slice(2, 4), 16), parseInt(h.slice(4, 6), 16), h.length >= 8 ? parseInt(h.slice(6, 8), 16) / 255 : 1];
42
+ }
43
+
44
+ class NadirMap {
45
+ // tokens: [{role, rgba:[r,g,b,a]}]; outputs: same length, each '#rrggbb[aa]' | [r,g,b,a] | null
46
+ // opts.keep: set of token indices that must stay unchanged (media overlays)
47
+ constructor(tokens, outputs, opts) {
48
+ opts = opts || {};
49
+ this.exact = {}; // role -> key -> css
50
+ this.pts = {}; // role -> [{lab, alpha, dl, da, db, dA}]
51
+ for (const r of ROLES) { this.exact[r] = new Map(); this.pts[r] = []; }
52
+ this.fillL = new Map(); // SVG fill colour -> output lightness
53
+ const keep = opts.keep || new Set();
54
+ tokens.forEach((t, i) => {
55
+ let o = outputs[i];
56
+ if (keep.has(i)) o = t.rgba;
57
+ if (o == null) return;
58
+ if (typeof o === 'string') o = parseHex(o);
59
+ const role = typeof t.role === 'number' ? ROLES[t.role] : t.role;
60
+ const [r, g, b, a] = t.rgba;
61
+ this.exact[role].set(keyOf(r, g, b, a), css([o[0], o[1], o[2]], o[3] == null ? a : o[3]));
62
+ const li = oklab(r, g, b), lo = oklab(o[0], o[1], o[2]);
63
+ if (role === 'fill' && !keep.has(i)) this.fillL.set(keyOf(r, g, b, a), lo[0]); // media keeps say nothing about ink
64
+ this.pts[role].push({lab: li, alpha: a, d: [lo[0] - li[0], lo[1] - li[1], lo[2] - li[2]], out: lo, oA: o[3] == null ? a : o[3]});
65
+ });
66
+ this.cache = new Map();
67
+ this.maxDist = opts.maxDist || 0.16;
68
+ }
69
+ // An SVG fill the model saw on this page: drawn as light ink (logo, icon) or as a dark surface (wave, panel)?
70
+ // null when the colour was never seen as a fill: the engine then falls back to its own size heuristic.
71
+ fillSide(r, g, b, a) {
72
+ const L = this.fillL.get(keyOf(r, g, b, a));
73
+ return L == null ? null : (L >= 0.5 ? 'fg' : 'bg');
74
+ }
75
+ lookup(role, r, g, b, a) {
76
+ const ck = role + keyOf(r, g, b, a);
77
+ let v = this.cache.get(ck);
78
+ if (v !== undefined) return v;
79
+ v = this._lookup(role, r, g, b, a);
80
+ this.cache.set(ck, v);
81
+ return v;
82
+ }
83
+ _lookup(role, r, g, b, a) {
84
+ const key = keyOf(r, g, b, a);
85
+ const chain = CHAIN[role] || [role];
86
+ for (const rr of chain) { const e = this.exact[rr].get(key); if (e) return e; }
87
+ // nearest-neighbour transfer within the role chain (first role with points)
88
+ const q = oklab(r, g, b);
89
+ let best = [];
90
+ for (let ci = 0; ci < chain.length; ci++) {
91
+ const pts = this.pts[chain[ci]];
92
+ for (const p of pts) {
93
+ const dl = p.lab[0] - q[0], da = p.lab[1] - q[1], db = p.lab[2] - q[2];
94
+ const d = Math.sqrt(dl * dl * 1.5 + da * da + db * db + (p.alpha - a) * (p.alpha - a) * 0.25) + ci * 0.03;
95
+ best.push([d, p]);
96
+ }
97
+ if (best.length >= 6) break;
98
+ }
99
+ if (!best.length) return null;
100
+ best.sort((x, y) => x[0] - y[0]);
101
+ best = best.slice(0, 4);
102
+ if (best[0][0] > this.maxDist) return null; // too far from anything the model saw: built-in fallback
103
+ let w = 0, d0 = 0, d1 = 0, d2 = 0, dA = 0;
104
+ for (const [d, p] of best) {
105
+ const wi = 1 / (d * d + 1e-4);
106
+ w += wi; d0 += p.d[0] * wi; d1 += p.d[1] * wi; d2 += p.d[2] * wi; dA += (p.oA - p.alpha) * wi;
107
+ }
108
+ const L = Math.max(0, Math.min(1, q[0] + d0 / w));
109
+ // keep chroma direction of the query, scale like the neighbours did
110
+ const rgb = toRgb(L, q[1] + d1 / w, q[2] + d2 / w);
111
+ const alpha = Math.max(0, Math.min(1, a + dA / w));
112
+ return css(rgb, alpha);
113
+ }
114
+ }
115
+
116
+ // image treatments as CSS
117
+ const TREAT_CSS = {
118
+ invert: 'filter: invert(1) hue-rotate(180deg) !important;',
119
+ light: 'filter: brightness(0) invert(0.92) !important;',
120
+ dim: 'filter: brightness(0.88) !important;',
121
+ 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;',
122
+ };
123
+ // Content images must never be colour-altered: only flat graphics (logos, icons, diagrams) may get a treatment.
124
+ // a logo on its own white rectangle: adapt it instead of leaving a white box. Only images the page marks as its
125
+ // logo (logo/brand/wordmark naming, or the link home); product shots, photos and other images are never touched
126
+ // Returns 'adapt:<fallback>': the page tries the adaptation (its pixel check confirms the white background)
127
+ // and uses the fallback treatment when there is nothing to adapt. Without pixel stats yet, logos still try.
128
+ function adaptable(im, s, surfL) {
129
+ if (!(im.logo || (im.homeLink && im.header)) || im.h > 200 || (surfL != null && surfL >= 0.4)) return false;
130
+ if (!s) return im.w * im.h <= 250000;
131
+ return s.transp < 0.05 && (s.ringA || 0) >= 0.95 && (s.ringL || 0) >= 0.85 && (s.white || 0) >= 0.3 && s.colors < 0.2;
132
+ }
133
+ function baseTreatment(t, im, surfL) {
134
+ const s = im && im.stats;
135
+ // a light, flat CSS background image (texture, pale wash, pattern) sits behind content whose text is now light:
136
+ // clear it so the dark surface designed for that element shows through. Photos and detailed images stay.
137
+ // (almost no lightness variation, so never a photo; SVG backgrounds can't be photos, so a smooth gradient qualifies).
138
+ // Logos drawn as backgrounds are never cleared.
139
+ if (im && im.kind === 'bg' && !im.logo && !im.homeLink && s && s.transp < 0.05 && s.meanL > 0.8 && im.w * im.h >= 40000 &&
140
+ s.stdL < (/\.svg$/i.test((im.src || '').split('?')[0]) ? 0.15 : 0.07)) return 'clear';
141
+ // a mostly dark graphic on transparency would vanish on a dark surface: hue-preserving inversion
142
+ 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';
143
+ if (!t || t === 'keep') return 'keep';
144
+ const photoLike = s ? (s.transp < 0.02 && s.colors > 0.18 && s.stdL > 0.12) : (im && im.w * im.h > 60000);
145
+ if (photoLike && t !== 'dim') return 'keep';
146
+ if (t === 'dim' && im && im.w * im.h < 20000) return 'keep';
147
+ // a transparent graphic that is already light reads fine on a dark surface: leave it alone
148
+ if ((t === 'plate' || t === 'light' || t === 'invert') && s && s.transp > 0.15 && s.meanL >= 0.55 && (s.black || 0) < 0.15 &&
149
+ (surfL == null || surfL < 0.4)) return 'keep';
150
+ // no light boxes on a dark page: a dark logo on transparency gets a hue-preserving inversion instead of a plate,
151
+ // an opaque one already carries its own background
152
+ if (t === 'plate') return (!s || s.transp > 0.15) ? 'invert' : 'keep';
153
+ return t;
154
+ }
155
+ function safeTreatment(t, im, surfL) {
156
+ const base = baseTreatment(t, im, surfL);
157
+ return base !== 'clear' && im && adaptable(im, im.stats, surfL) ? 'adapt:' + base : base;
158
+ }
159
+
160
+ // ---------- flat graphics on a white background (logos, icons, diagrams): adapt the image itself ----------
161
+ // The white background becomes transparent so the page's own dark surface shows through, neutral ink flips to
162
+ // light, and brand colours keep their hue and are lifted only as far as needed to read on dark.
163
+ function adaptPixels(d, w, h, info) {
164
+ info = info || {};
165
+ let n = 0, light = 0, sr = 0, sg = 0, sb = 0;
166
+ const ringPx = (p) => {
167
+ const i = p * 4;
168
+ if (d[i + 3] < 250) return;
169
+ n++;
170
+ const o = oklab(d[i], d[i + 1], d[i + 2]);
171
+ if (o[0] > 0.9 && Math.hypot(o[1], o[2]) < 0.04) { light++; sr += d[i]; sg += d[i + 1]; sb += d[i + 2]; }
172
+ };
173
+ for (let x = 0; x < w; x++) { ringPx(x); ringPx((h - 1) * w + x); }
174
+ for (let y = 1; y < h - 1; y++) { ringPx(y * w); ringPx(y * w + w - 1); }
175
+ info.ring = n ? light / n : 0;
176
+ if (!n || light / n < 0.5) { info.why = 'ring'; return false; } // no light background to remove (tight crops touch the edge)
177
+ const bg = oklab(sr / light, sg / light, sb / light);
178
+ const cache = new Map();
179
+ for (let p = 0, P = w * h; p < P; p++) {
180
+ const i = p * 4;
181
+ if (!d[i + 3]) continue;
182
+ const key = (d[i] << 16) | (d[i + 1] << 8) | d[i + 2];
183
+ let r = cache.get(key);
184
+ if (!r) {
185
+ const o = oklab(d[i], d[i + 1], d[i + 2]);
186
+ const dist = Math.hypot(o[0] - bg[0], o[1] - bg[1], o[2] - bg[2]);
187
+ const cov = Math.max(0, Math.min(1, (dist - 0.03) / 0.22)); // 0 = background, 1 = ink, between = anti-aliased edge
188
+ if (cov <= 0) r = [0, 0, 0, 0];
189
+ else {
190
+ // unmix the edge pixel from the background, then restyle the ink for a dark surface
191
+ 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;
192
+ L = Math.max(0, Math.min(1, L));
193
+ const C = Math.hypot(a, b);
194
+ if (C < 0.05) L = 0.94 - 0.75 * L; // black/gray ink -> light
195
+ else if (L < 0.62) { L = 0.62 + 0.15 * L / 0.62; a *= 0.92; b *= 0.92; } // brand colour: same hue, lifted to read on dark
196
+ const rgb = toRgb(L, a, b);
197
+ r = [rgb[0], rgb[1], rgb[2], cov];
198
+ }
199
+ cache.set(key, r);
200
+ }
201
+ d[i] = r[0]; d[i + 1] = r[1]; d[i + 2] = r[2]; d[i + 3] = Math.round(d[i + 3] * r[3]);
202
+ }
203
+ return true;
204
+ }
205
+ // blob -> data URL of the adapted image at its natural size (so layout never shifts), or null
206
+ async function sniffType(blob) {
207
+ const h = new Uint8Array(await blob.slice(0, 64).arrayBuffer());
208
+ const txt = String.fromCharCode(...h);
209
+ if (h[0] === 0x89 && txt.slice(1, 4) === 'PNG') return 'image/png';
210
+ if (h[0] === 0xff && h[1] === 0xd8) return 'image/jpeg';
211
+ if (txt.startsWith('GIF8')) return 'image/gif';
212
+ if (txt.startsWith('RIFF') && txt.slice(8, 12) === 'WEBP') return 'image/webp';
213
+ if (/<svg|<\?xml/i.test(txt)) return 'image/svg+xml';
214
+ return blob.type;
215
+ }
216
+ async function adaptImage(blob, info) {
217
+ info = info || {};
218
+ const type = await sniffType(blob); // servers mislabel images (e.g. video/webp)
219
+ if (type !== blob.type) blob = new Blob([blob], {type});
220
+ let src = null;
221
+ if (!/svg/.test(blob.type)) { try { src = await createImageBitmap(blob); } catch (e) { src = null; } }
222
+ if (!src) {
223
+ const u = URL.createObjectURL(blob.type ? blob : new Blob([blob], {type: 'image/svg+xml'}));
224
+ src = await new Promise((res) => { const im = new Image(); im.onload = () => res(im); im.onerror = () => res(null); im.src = u; });
225
+ URL.revokeObjectURL(u);
226
+ }
227
+ if (!src) { info.why = 'decode'; return null; }
228
+ const w = src.naturalWidth || src.width, h = src.naturalHeight || src.height;
229
+ if (!w || !h || w * h > 2400 * 2400) { info.why = 'size'; return null; }
230
+ const c = document.createElement('canvas'); c.width = w; c.height = h;
231
+ const ctx = c.getContext('2d', {willReadFrequently: true});
232
+ ctx.drawImage(src, 0, 0, w, h);
233
+ if (src.close) src.close();
234
+ const im = ctx.getImageData(0, 0, w, h);
235
+ if (!adaptPixels(im.data, w, h, info)) return null;
236
+ ctx.putImageData(im, 0, 0);
237
+ return c.toDataURL('image/png');
238
+ }
239
+
240
+ globalThis.NadirMap = NadirMap;
241
+ globalThis.NadirAdaptImage = adaptImage;
242
+ globalThis.NadirImageCSS = TREAT_CSS;
243
+ globalThis.NadirSafeTreatment = safeTreatment;
244
+ })();
training/build_ds.py ADDED
@@ -0,0 +1,281 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """Build Nadir training data: page features + per-token dark targets (Haiku designs and native dark-mode pairs)."""
3
+ import argparse, collections, difflib, gzip, hashlib, json, math, os, random, sys
4
+ from pathlib import Path
5
+ import numpy as np
6
+
7
+ import safety
8
+
9
+ ROLES = ["bg", "text", "border", "shadow", "gradient", "fill", "stroke"]
10
+ TREAT = ["keep", "dim", "invert", "light", "plate"]
11
+ MAX_TOK, MAX_IMG = 256, 32
12
+
13
+
14
+ def lin(u):
15
+ u = np.asarray(u, dtype=np.float64) / 255.0
16
+ return np.where(u <= 0.04045, u / 12.92, ((u + 0.055) / 1.055) ** 2.4)
17
+
18
+
19
+ def oklab(rgb):
20
+ r, g, b = lin(rgb[..., 0]), lin(rgb[..., 1]), lin(rgb[..., 2])
21
+ l = np.cbrt(0.4122214708 * r + 0.5363325363 * g + 0.0514459929 * b)
22
+ m = np.cbrt(0.2119034982 * r + 0.6806995451 * g + 0.1073969566 * b)
23
+ s = np.cbrt(0.0883024619 * r + 0.2817188376 * g + 0.6299787005 * b)
24
+ return np.stack([0.2104542553 * l + 0.7936177850 * m - 0.0040720468 * s,
25
+ 1.9779984951 * l - 2.4285922050 * m + 0.4505937099 * s,
26
+ 0.0259040371 * l + 0.7827717662 * m - 0.8086757660 * s], -1)
27
+
28
+
29
+ def rel_lum(rgb):
30
+ c = lin(rgb)
31
+ return 0.2126 * c[..., 0] + 0.7152 * c[..., 1] + 0.0722 * c[..., 2]
32
+
33
+
34
+ def hex2rgba(h):
35
+ h = h.lstrip("#")
36
+ return [int(h[0:2], 16), int(h[2:4], 16), int(h[4:6], 16), int(h[6:8], 16) / 255 if len(h) >= 8 else 1.0]
37
+
38
+
39
+ def contrast(fg, bg):
40
+ l1, l2 = rel_lum(np.array(fg[:3])), rel_lum(np.array(bg[:3]))
41
+ return (max(l1, l2) + 0.05) / (min(l1, l2) + 0.05)
42
+
43
+
44
+ def over(top, bottom):
45
+ a = top[3]
46
+ return [top[i] * a + bottom[i] * (1 - a) for i in range(3)] + [1.0]
47
+
48
+
49
+ def page_arrays(L):
50
+ toks = L["tokens"][:MAX_TOK]
51
+ n = len(toks)
52
+ F = np.array([t["f"] for t in toks], dtype=np.float32)
53
+ role = np.array([t["role"] for t in toks], dtype=np.int64)
54
+ rgba = np.array([t["rgba"] for t in toks], dtype=np.float32)
55
+ rel = np.array([r for r in L["rel"] if r[1] < n and r[2] < n], dtype=np.float32).reshape(-1, 4)
56
+ imgs = L["images"][:MAX_IMG]
57
+ IF = np.array([im["f"] for im in imgs], dtype=np.float32).reshape(len(imgs), 39)
58
+ thumbs = np.zeros((len(imgs), 32, 32, 4), dtype=np.uint8)
59
+ for j, im in enumerate(imgs):
60
+ if im.get("thumb") and len(im["thumb"]) == 4096:
61
+ thumbs[j] = np.array(im["thumb"], dtype=np.uint8).reshape(32, 32, 4)
62
+ img_under = np.array([im.get("underTok", -1) if (im.get("underTok", -1) < n) else -1 for im in imgs], dtype=np.int64)
63
+ w = []
64
+ for t in toks:
65
+ d = t.get("d") or {}
66
+ w.append(1.0 + math.log1p(d.get("area", 0) / 2000) + math.log1p(d.get("chars", 0) / 20) + 0.5 * math.log1p(d.get("nVis", 0)) + 0.3 * math.log1p(d.get("rules", 0)))
67
+ keep = [i for i, t in enumerate(toks) if safety.is_media_keep(t)]
68
+ page = np.array([L["vw"] / 1280, min(8, L["docH"] / L["vh"]) / 8, n / MAX_TOK, len(imgs) / MAX_IMG, L["darkness"]["bgL"], L["darkness"]["textL"],
69
+ 1.0 if L.get("canvasDark") else 0.0, min(1, L["stats"].get("tokensDropped", 0) / 200)], dtype=np.float32)
70
+ F, IF, page = (np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0) for x in (F, IF, page))
71
+ return dict(F=F, role=role, rgba=rgba, rel=rel, IF=IF, thumbs=thumbs, img_under=img_under, w=np.array(w, dtype=np.float32), page=page, keep=np.array(keep, dtype=np.int64))
72
+
73
+
74
+ def haiku_targets(L, h, arrs, qc):
75
+ n = len(arrs["role"])
76
+ tgt = np.full((n, 4), np.nan, dtype=np.float32)
77
+ for k, v in h["colors"].items():
78
+ i = int(k)
79
+ if i < n:
80
+ tgt[i] = hex2rgba(v)
81
+ imgt = np.full((len(arrs["IF"]),), -1, dtype=np.int64)
82
+ for k, v in (h.get("images") or {}).items():
83
+ j = int(k)
84
+ if j < len(imgt) and v in TREAT:
85
+ imgt[j] = TREAT.index(v)
86
+ # ---- QC ----
87
+ asked = min(n, h.get("n_tok", n))
88
+ got = np.isfinite(tgt[:asked, 0]).sum()
89
+ if asked == 0 or got / asked < 0.9:
90
+ qc["coverage"] += 1
91
+ return None
92
+ root = [i for i, t in enumerate(L["tokens"][:n]) if (t.get("d") or {}).get("root")]
93
+ if root and np.isfinite(tgt[root[0], 0]):
94
+ rl = oklab(np.array(tgt[root[0], :3]))[0]
95
+ if rl > 0.4:
96
+ qc["light_root"] += 1
97
+ return None
98
+ # text contrast vs mapped dominant surface
99
+ tb = collections.defaultdict(list)
100
+ for typ, a, b, w in arrs["rel"]:
101
+ if int(typ) == 0:
102
+ tb[int(a)].append((w, int(b)))
103
+ bad = tot = 0.0
104
+ for i, t in enumerate(L["tokens"][:n]):
105
+ if t["role"] != 1 or not np.isfinite(tgt[i, 0]):
106
+ continue
107
+ chars = (t.get("d") or {}).get("chars", 0)
108
+ if not chars or not tb.get(i):
109
+ continue
110
+ j = max(tb[i])[1]
111
+ if not np.isfinite(tgt[j, 0]):
112
+ continue
113
+ bgc = tgt[j]
114
+ base = over(list(bgc), [18, 18, 18, 1]) if bgc[3] < 1 else list(bgc)
115
+ fg = over(list(tgt[i]), base)
116
+ c = contrast(fg, base)
117
+ tot += chars
118
+ if c < 3.0:
119
+ bad += chars
120
+ if tot and bad / tot > 0.15:
121
+ qc["contrast"] += 1
122
+ return None
123
+ qc["ok"] += 1
124
+ return tgt, imgt
125
+
126
+
127
+ def repair_contrast(arrs, tgt):
128
+ """Apply the runtime rules (hue guard, invisible-stays-invisible, contrast) to a Haiku design, so the model
129
+ learns what the runtime would show."""
130
+ n = len(arrs["role"])
131
+ toks = [{"role": int(arrs["role"][i]), "rgba": arrs["rgba"][i].tolist(), "f": arrs["F"][i]} for i in range(n)]
132
+ out, keep = [], set(int(k) for k in arrs["keep"])
133
+ for i in range(n):
134
+ if not np.isfinite(tgt[i, 0]):
135
+ keep.add(i)
136
+ out.append({"rgb": [0, 0, 0], "a": 0.0, "L": 0.0})
137
+ continue
138
+ rgb = [int(round(float(x))) for x in tgt[i, :3]]
139
+ out.append({"rgb": rgb, "a": float(tgt[i, 3]), "L": safety.oklab(*rgb)[0]})
140
+ before = [list(o["rgb"]) for o in out]
141
+ for i in range(n):
142
+ if i not in keep:
143
+ safety.hue_guard(toks[i], out[i])
144
+ safety.keep_invisible(toks, out, arrs["rel"].tolist(), keep)
145
+ safety.enforce_contrast(toks, out, arrs["rel"].tolist(), keep)
146
+ changed = 0
147
+ for i in range(n):
148
+ if out[i]["rgb"] != before[i]:
149
+ tgt[i, :3] = out[i]["rgb"]
150
+ changed += 1
151
+ return changed
152
+
153
+
154
+ def native_targets(L, D, arrs, qc):
155
+ le, de = L.get("ec") or [], D.get("ec") or []
156
+ if not le or not de:
157
+ qc["no_ec"] += 1
158
+ return None
159
+ if len(le) == len(de) and all(a[0] == b[0] for a, b in zip(le, de)):
160
+ pairs = list(zip(range(len(le)), range(len(de))))
161
+ else:
162
+ sm = difflib.SequenceMatcher(None, [r[0] for r in le], [r[0] for r in de], autojunk=False)
163
+ pairs = []
164
+ for blk in sm.get_matching_blocks():
165
+ pairs += [(blk.a + k, blk.b + k) for k in range(blk.size)]
166
+ if len(pairs) < 0.8 * len(le):
167
+ qc["align"] += 1
168
+ return None
169
+ votes = collections.defaultdict(collections.Counter)
170
+ for ia, ib in pairs:
171
+ a, b = le[ia], de[ib]
172
+ w = (a[1] or 0) + 1
173
+ if a[2]:
174
+ # bg: effective dark surface at this element (opaque light bg) or raw dark bg (translucent)
175
+ alpha = int(a[2][6:8], 16)
176
+ tgt = (b[8] if alpha >= 250 else b[2]) or b[8]
177
+ votes[(0, a[2])][tgt] += w
178
+ for j, role in ((3, 1), (4, 2), (5, 5), (6, 6)):
179
+ if a[j]:
180
+ votes[(role, a[j])][b[j] or "00000000"] += w
181
+ if a[7] and b[7]:
182
+ la, lb = a[7].split(","), b[7].split(",")
183
+ if len(la) == len(lb):
184
+ for x, y in zip(la, lb):
185
+ votes[(4, x)][y] += w
186
+ n = len(arrs["role"])
187
+ tgt = np.full((n, 4), np.nan, dtype=np.float32)
188
+ hit = 0
189
+ for i, t in enumerate(L["tokens"][:n]):
190
+ key = t["k"].split(":")[1]
191
+ c = votes.get((t["role"], key))
192
+ if not c:
193
+ continue
194
+ best, bw = c.most_common(1)[0]
195
+ if bw / sum(c.values()) < 0.5:
196
+ continue
197
+ tgt[i] = hex2rgba(best)
198
+ hit += 1
199
+ if hit < 4:
200
+ qc["few"] += 1
201
+ return None
202
+ qc["ok"] += 1
203
+ return tgt, np.full((len(arrs["IF"]),), -1, dtype=np.int64)
204
+
205
+
206
+ def main():
207
+ ap = argparse.ArgumentParser()
208
+ ap.add_argument("--crawl", required=True)
209
+ ap.add_argument("--haiku", nargs="*", default=[])
210
+ ap.add_argument("--out", required=True)
211
+ ap.add_argument("--val-frac", type=float, default=0.05)
212
+ args = ap.parse_args()
213
+ crawl = Path(args.crawl)
214
+ index = {}
215
+ for line in (l for ip in sorted(crawl.glob("index*.jsonl")) for l in open(ip)):
216
+ try:
217
+ r = json.loads(line)
218
+ index[r["id"]] = r
219
+ except Exception:
220
+ pass
221
+ samples = []
222
+ qc_h = collections.Counter(); qc_n = collections.Counter()
223
+ # Haiku-designed pages (source 0)
224
+ for hd in args.haiku:
225
+ for fp in Path(hd).glob("*.json"):
226
+ if fp.name in ("ledger.jsonl",):
227
+ continue
228
+ try:
229
+ h = json.loads(fp.read_text())
230
+ except Exception:
231
+ continue
232
+ if h.get("refusal") or "colors" not in h:
233
+ continue
234
+ pid = h["id"]
235
+ try:
236
+ L = json.load(gzip.open(crawl / "pages" / pid / "light.json.gz"))
237
+ except Exception:
238
+ continue
239
+ arrs = page_arrays(L)
240
+ r = haiku_targets(L, h, arrs, qc_h)
241
+ if r is None:
242
+ continue
243
+ arrs["tgt"], arrs["imgt"] = r
244
+ if repair_contrast(arrs, arrs["tgt"]):
245
+ qc_h["repaired"] += 1
246
+ arrs["src"] = 0
247
+ arrs["id"] = pid
248
+ arrs["host"] = index.get(pid, {}).get("url", L["url"]).split("/")[2]
249
+ samples.append(arrs)
250
+ # native dark-mode pairs (source 1)
251
+ for pid, r in index.items():
252
+ if r.get("kind") not in ("native_media", "native_class"):
253
+ continue
254
+ try:
255
+ L = json.load(gzip.open(crawl / "pages" / pid / "light.json.gz"))
256
+ D = json.load(gzip.open(crawl / "pages" / pid / "dark.json.gz"))
257
+ except Exception:
258
+ continue
259
+ arrs = page_arrays(L)
260
+ res = native_targets(L, D, arrs, qc_n)
261
+ if res is None:
262
+ continue
263
+ arrs["tgt"], arrs["imgt"] = res
264
+ arrs["src"] = 1
265
+ arrs["id"] = pid
266
+ arrs["host"] = r["url"].split("/")[2]
267
+ samples.append(arrs)
268
+ print("haiku QC", dict(qc_h), "native QC", dict(qc_n), "samples", len(samples), flush=True)
269
+ hosts = sorted(set(s["host"].replace("www.", "") for s in samples))
270
+ # stable split: a host's side never changes as the crawl grows
271
+ val_hosts = {h for h in hosts if int(hashlib.md5(h.encode()).hexdigest()[:8], 16) % 1000 < args.val_frac * 1000}
272
+ train = [s for s in samples if s["host"].replace("www.", "") not in val_hosts]
273
+ val = [s for s in samples if s["host"].replace("www.", "") in val_hosts]
274
+ import pickle
275
+ with open(args.out, "wb") as f:
276
+ pickle.dump({"train": train, "val": val, "roles": ROLES, "treat": TREAT}, f, protocol=4)
277
+ 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))
278
+
279
+
280
+ if __name__ == "__main__":
281
+ main()
training/export.py ADDED
@@ -0,0 +1,211 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """Export NadirNet to ONNX (fp32/fp16) and build calibrated 8-bit and 4-bit variants; measure each against fp32."""
3
+ import argparse, json, os, pickle, time
4
+ import numpy as np
5
+ import torch
6
+ import torch.nn as nn
7
+
8
+ from model import NadirNet
9
+ from train import collate, prep, evaluate
10
+
11
+ NAMES_IN = ["tok_f", "tok_role", "tok_lab", "img_f", "img_thumb", "page_f", "src", "rel", "mask"]
12
+ NAMES_OUT = ["lab", "alpha_logit", "img_logits"]
13
+
14
+
15
+ class Wrapped(nn.Module):
16
+ def __init__(self, m):
17
+ super().__init__()
18
+ self.m = m
19
+
20
+ def forward(self, tok_f, tok_role, tok_lab, img_f, img_thumb, page_f, src, rel, mask):
21
+ return self.m(tok_f, tok_role, tok_lab, img_f, img_thumb, page_f, src, rel, mask)
22
+
23
+
24
+ def feeds_for(s, pad=None):
25
+ b = collate([s], "cpu")
26
+ if pad: # fixed shapes for calibration (quantizers stack activations across samples)
27
+ N0, M0 = b["tok_f"].shape[1], b["img_f"].shape[1]
28
+ N, M = pad
29
+ if N0 <= N and M0 <= M:
30
+ T0, T = 1 + N0 + M0, 1 + N + M
31
+ def padto(x, dims, val=0):
32
+ shape = list(x.shape)
33
+ for d, sz in dims.items():
34
+ shape[d] = sz
35
+ y = torch.full(shape, val, dtype=x.dtype)
36
+ y[tuple(slice(0, x.shape[i]) for i in range(x.ndim))] = x
37
+ return y
38
+ idx = list(range(1 + N0)) + list(range(1 + N, 1 + N + M0))
39
+ rel = torch.zeros(1, T, T, b["rel"].shape[-1]); mask = torch.zeros(1, T, dtype=torch.bool)
40
+ ii = torch.tensor(idx)
41
+ rel[0][ii[:, None], ii[None, :]] = b["rel"][0]
42
+ mask[0, ii] = b["mask"][0]
43
+ b = dict(b, tok_f=padto(b["tok_f"], {1: N}), role=padto(b["role"], {1: N}), lab=padto(b["lab"], {1: N}),
44
+ img_f=padto(b["img_f"], {1: M}), thumbs=padto(b["thumbs"], {1: M}), rel=rel, mask=mask)
45
+ M = b["img_f"].shape[1]
46
+ if M == 0: # the runtime always sends at least one (masked) image slot
47
+ T = b["mask"].shape[1]
48
+ b["img_f"] = torch.zeros(1, 1, 39); b["thumbs"] = torch.zeros(1, 1, 4, 32, 32)
49
+ rel = torch.zeros(1, T + 1, T + 1, b["rel"].shape[-1]); rel[:, :T, :T] = b["rel"]; b["rel"] = rel
50
+ b["mask"] = torch.cat([b["mask"], torch.zeros(1, 1, dtype=torch.bool)], 1)
51
+ return {"tok_f": b["tok_f"].numpy(), "tok_role": b["role"].numpy(), "tok_lab": b["lab"].numpy(), "img_f": b["img_f"].numpy(),
52
+ "img_thumb": b["thumbs"].numpy(), "page_f": b["page"].numpy(), "src": b["src"].numpy(), "rel": b["rel"].numpy(), "mask": b["mask"].numpy()}
53
+
54
+
55
+ def fidelity(sess, ref_sess, samples):
56
+ """Mean OKLab distance between this model's colours and fp32's, plus image-decision agreement."""
57
+ de, agree, n_img, worst = [], 0, 0, 0.0
58
+ for s in samples:
59
+ f = feeds_for(s)
60
+ a = sess.run(None, f); r = ref_sess.run(None, f)
61
+ n = s["F"].shape[0]
62
+ d = np.linalg.norm(a[0][0, :n] - r[0][0, :n], axis=-1)
63
+ w = s["w"][:n]
64
+ de.append(float((d * w).sum() / w.sum())); worst = max(worst, float(d.max()))
65
+ m = len(s["IF"])
66
+ if m:
67
+ agree += int((a[2][0, :m].argmax(-1) == r[2][0, :m].argmax(-1)).sum()); n_img += m
68
+ return {"dE_vs_fp32": float(np.mean(de)), "dE_max": worst, "img_agree": agree / max(1, n_img)}
69
+
70
+
71
+ def bench(sess, samples, reps=3):
72
+ ts = []
73
+ for s in samples[:20]:
74
+ f = feeds_for(s)
75
+ sess.run(None, f)
76
+ t = time.time()
77
+ for _ in range(reps):
78
+ sess.run(None, f)
79
+ ts.append((time.time() - t) / reps * 1000)
80
+ return float(np.median(ts))
81
+
82
+
83
+ def main():
84
+ ap = argparse.ArgumentParser()
85
+ ap.add_argument("--ckpt", required=True)
86
+ ap.add_argument("--data", required=True)
87
+ ap.add_argument("--out", required=True)
88
+ ap.add_argument("--calib", type=int, default=256)
89
+ ap.add_argument("--only-float", action="store_true", help="fp32/fp16 only (calibrated 8/4-bit come from quantize.py)")
90
+ args = ap.parse_args()
91
+ os.makedirs(args.out, exist_ok=True)
92
+ import onnx, onnxruntime as ort
93
+ ck = torch.load(args.ckpt, map_location="cpu")
94
+ model = NadirNet(**{k: v for k, v in ck["cfg"].items()}); model.load_state_dict(ck["model"]); model.eval()
95
+ data = pickle.load(open(args.data, "rb"))
96
+ val = [prep(s) for s in data["val"]]
97
+ rng = np.random.default_rng(0)
98
+ calib = [prep(s) for s in rng.choice(np.array(data["train"], dtype=object), size=min(args.calib, len(data["train"])), replace=False)]
99
+ s0 = max(val, key=lambda s: len(s["role"]))
100
+ ex = feeds_for(s0)
101
+ par_samples = val[:8]
102
+ with torch.no_grad():
103
+ ref_out = [model(*[torch.from_numpy(feeds_for(s)[k]) for k in NAMES_IN])[0].numpy() for s in par_samples]
104
+ fp32 = os.path.join(args.out, "nadir-fp32.onnx")
105
+ 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,
106
+ dynamic_axes={"tok_f": {1: "N"}, "tok_role": {1: "N"}, "tok_lab": {1: "N"}, "img_f": {1: "M"}, "img_thumb": {1: "M"},
107
+ "rel": {1: "T", 2: "T"}, "mask": {1: "T"}, "lab": {1: "N"}, "alpha_logit": {1: "N"}, "img_logits": {1: "M"}},
108
+ opset_version=17, do_constant_folding=True, dynamo=False)
109
+ onnx.checker.check_model(fp32)
110
+ so = ort.SessionOptions(); so.intra_op_num_threads = 4
111
+ ref = ort.InferenceSession(fp32, so, providers=["CPUExecutionProvider"])
112
+ # parity torch vs onnx (reference computed before export)
113
+ parity = max(float(np.abs(r - ref.run(None, feeds_for(s))[0]).max()) for r, s in zip(ref_out, par_samples))
114
+ report = {"params": sum(p.numel() for p in model.parameters()), "parity_max_abs": parity, "variants": {}}
115
+ print("parity", parity, flush=True)
116
+
117
+ # fp16 (weights + compute in fp16, IO kept fp32)
118
+ # fp16: export a half-precision copy with fp32 inputs/outputs (casts inside the graph)
119
+ have16 = True
120
+ try:
121
+ class Half(nn.Module):
122
+ def __init__(self, m):
123
+ super().__init__(); self.m = m.half()
124
+ def forward(self, tok_f, tok_role, tok_lab, img_f, img_thumb, page_f, src, rel, mask):
125
+ o = self.m(tok_f.half(), tok_role, tok_lab.half(), img_f.half(), img_thumb.half(), page_f.half(), src, rel.half(), mask)
126
+ return tuple(x.float() for x in o)
127
+ import copy
128
+ 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"),
129
+ input_names=NAMES_IN, output_names=NAMES_OUT,
130
+ dynamic_axes={"tok_f": {1: "N"}, "tok_role": {1: "N"}, "tok_lab": {1: "N"}, "img_f": {1: "M"}, "img_thumb": {1: "M"},
131
+ "rel": {1: "T", 2: "T"}, "mask": {1: "T"}, "lab": {1: "N"}, "alpha_logit": {1: "N"}, "img_logits": {1: "M"}},
132
+ opset_version=17, do_constant_folding=True, dynamo=False)
133
+ ort.InferenceSession(os.path.join(args.out, "nadir-fp16.onnx"), providers=["CPUExecutionProvider"])
134
+ except Exception as e:
135
+ print("fp16 failed:", str(e)[:300], flush=True); have16 = False
136
+
137
+ if args.only_float:
138
+ variants = {"fp32": fp32}
139
+ if have16: variants["fp16"] = os.path.join(args.out, "nadir-fp16.onnx")
140
+ for name, path in variants.items():
141
+ sess = ort.InferenceSession(path, so, providers=["CPUExecutionProvider"])
142
+ 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)
143
+ report["variants"][name] = r
144
+ print(name, r, flush=True)
145
+ json.dump(report, open(os.path.join(args.out, "export_report.json"), "w"), indent=1)
146
+ return
147
+ from onnxruntime.quantization import quantize_dynamic, quantize_static, QuantType, QuantFormat, CalibrationDataReader, CalibrationMethod
148
+ from onnxruntime.quantization.shape_inference import quant_pre_process
149
+ pre = os.path.join(args.out, "_pre.onnx")
150
+ quant_pre_process(fp32, pre, skip_symbolic_shape=True)
151
+
152
+ class Reader(CalibrationDataReader):
153
+ def __init__(self, ss):
154
+ self.it = iter([feeds_for(s, pad=(256, 32)) for s in ss])
155
+ def get_next(self):
156
+ return next(self.it, None)
157
+
158
+ # 8-bit A: per-channel int8 weights, dynamic activations
159
+ q8d = os.path.join(args.out, "nadir-int8-dynamic.onnx")
160
+ quantize_dynamic(pre, q8d, weight_type=QuantType.QInt8, per_channel=True, op_types_to_quantize=["MatMul", "Gemm"])
161
+ # 8-bit B: static QDQ, activations calibrated on real pages (percentile)
162
+ q8s = os.path.join(args.out, "nadir-int8-static.onnx")
163
+ try:
164
+ quantize_static(pre, q8s, Reader(calib[:128]), quant_format=QuantFormat.QDQ, per_channel=True, activation_type=QuantType.QInt8,
165
+ weight_type=QuantType.QInt8, op_types_to_quantize=["MatMul", "Gemm"], calibrate_method=CalibrationMethod.Percentile,
166
+ extra_options={"CalibPercentile": 99.99, "ActivationSymmetric": True, "WeightSymmetric": True})
167
+ except Exception as e:
168
+ print("static int8 failed:", e, flush=True); q8s = None
169
+ # 4-bit: block-wise weights (MatMulNBits), GPTQ calibrated on real pages, RTN fallback
170
+ q4 = os.path.join(args.out, "nadir-int4.onnx")
171
+ from onnxruntime.quantization import matmul_nbits_quantizer as mnq
172
+ done4 = False
173
+ # GPTQ needs reduction dims that are multiples of the block size; leave tiny matrices (relation bias, 11 wide) in fp32
174
+ pm = onnx.load(pre)
175
+ inits = {i.name: i for i in pm.graph.initializer}
176
+ 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]
177
+ report["int4_fp32_nodes"] = len(skip)
178
+ for algo in ("gptq", "hqq", "rtn"):
179
+ try:
180
+ if algo == "gptq":
181
+ cfg = mnq.GPTQWeightOnlyQuantConfig(calibration_data_reader=Reader(calib), percdamp=0.01, block_size=32, actorder=True, mse=True)
182
+ elif algo == "hqq":
183
+ cfg = mnq.HQQWeightOnlyQuantConfig(block_size=32, bits=4)
184
+ else:
185
+ cfg = mnq.RTNWeightOnlyQuantConfig()
186
+ qz = mnq.MatMulNBitsQuantizer(onnx.load(pre), block_size=32, is_symmetric=False, accuracy_level=4, algo_config=cfg, nodes_to_exclude=skip)
187
+ qz.process()
188
+ qz.model.save_model_to_file(q4, use_external_data_format=False)
189
+ report["int4_algo"] = algo
190
+ done4 = True
191
+ break
192
+ except Exception as e:
193
+ print(f"int4 {algo} failed: {str(e)[:300]}", flush=True)
194
+ variants = {"fp32": fp32, "int8-dynamic": q8d}
195
+ if have16: variants["fp16"] = os.path.join(args.out, "nadir-fp16.onnx")
196
+ if q8s: variants["int8-static"] = q8s
197
+ if done4: variants["int4"] = q4
198
+ evs = val[:300]
199
+ for name, path in variants.items():
200
+ sess = ort.InferenceSession(path, so, providers=["CPUExecutionProvider"])
201
+ r = fidelity(sess, ref, evs)
202
+ r["size_mb"] = round(os.path.getsize(path) / 1e6, 2)
203
+ r["cpu_ms_median"] = round(bench(sess, evs), 2)
204
+ report["variants"][name] = r
205
+ print(name, r, flush=True)
206
+ os.remove(pre)
207
+ json.dump(report, open(os.path.join(args.out, "export_report.json"), "w"), indent=1)
208
+
209
+
210
+ if __name__ == "__main__":
211
+ main()
training/haiku_gen.py ADDED
@@ -0,0 +1,621 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """Nadir teacher data: Claude Haiku 5.5 designs a dark theme for each crawled light page.
3
+
4
+ Runs on Vega's host (the API key vault proxy is only reachable there). Writes one JSON per page to
5
+ OUT/<id>.json and keeps a running cost ledger; stops at --budget dollars.
6
+ """
7
+ import argparse, asyncio, base64, gzip, io, json, math, os, random, re, sys, time
8
+ from pathlib import Path
9
+
10
+ import anthropic
11
+ from PIL import Image, ImageDraw, ImageFont
12
+
13
+ MODEL = "claude-haiku-5-5"
14
+ PRICE_IN, PRICE_OUT = 0.10 / 1e6, 0.50 / 1e6 # prompts <= 100K tokens
15
+ PRICE_CACHE_READ, PRICE_CACHE_WRITE = 0.01 / 1e6, 0.125 / 1e6
16
+ ROLES = ["bg", "text", "border", "shadow", "gradient", "fill", "stroke"]
17
+ TREATMENTS = ["keep", "dim", "invert", "light", "plate"]
18
+
19
+ SYSTEM = """You are a senior product designer. A company has a light-themed website and has asked you to design its dark theme: the version their own design team would ship, not a mechanical inversion. You are given a screenshot of the light page, a contact sheet of its images, and the complete inventory of colors the page uses, each with its role and where it appears. You return the dark-theme replacement for every color and a treatment for every image.
20
+
21
+ How to design it:
22
+ - Surfaces. Choose one base dark surface for the page canvas. Near-black neutrals in the #0d0f12-#1c1e22 range work for most sites; a faint tint of the brand hue is welcome when it suits the brand. Avoid pure #000000 unless the brand is starkly monochrome. Keep the surface hierarchy of the light design but express it in dark-UI terms: anything that sits ON another surface (cards, panels, popovers, inputs, sticky headers, table stripes) is slightly LIGHTER than what is beneath it, in small steps (about 3-7 points of lightness). Alternating light-gray section bands become subtly different dark tones. Never let a nested surface end up darker than its parent unless it is a deliberate inset (code wells, input fields may be a touch darker).
23
+ - Colors that are already dark in the light design (dark heroes, footers, banners, code blocks, dark buttons) stay dark. Harmonize them with the new palette; never flip them to light.
24
+ - Text. Primary text is light but not pure white for body copy (around #e4e6ea-#f0f1f3). Secondary and muted text keep their hierarchy as progressively dimmer grays, and every text color must keep at least 4.5:1 contrast on the surface it actually sits on (3:1 for large or bold headings). Text that sits on a colored or dark-in-light-mode surface keeps working on that surface's new color.
25
+ - Brand and accent colors keep their hue so the brand stays recognizable. When an accent is used as text, a link or an icon on a dark surface, raise its lightness and slightly reduce chroma until it reads comfortably. When an accent is a filled button or badge background, keep it close to the original (deepen slightly if very light) and keep its label readable.
26
+ - Status colors (error, warning, success, info) keep their meaning and become readable on dark.
27
+ - Borders and dividers become subtle lines slightly lighter than the surface they sit on. Focus rings stay clearly visible.
28
+ - Translucent colors (8-digit hex): keep a similar alpha unless it no longer works. Black-tinted overlays used to darken or separate things in light mode usually become white-tinted at a similar or slightly lower alpha (rgba(0,0,0,.05) hover -> #ffffff0f). Shadows and scrims stay dark; shadows may get a little stronger.
29
+ - Hover, focus and active colors (marked as rule-only / state) keep their relationship to the resting state: if hover was darker in light mode, make it visibly lighter (or more emphasized) in dark mode.
30
+ - Overlays drawn on top of videos or images (video player controls, captions on photos, scrims, gradient fades over media) and the black backgrounds of media players are part of the media: return them unchanged.
31
+ - Syntax highlighting and data-visualization colors stay distinguishable from each other and readable on dark.
32
+ - The same input color in the same role always gets the same output. Every listed id must appear in your answer exactly once.
33
+
34
+ Images (choose one treatment each; the contact sheet shows transparency as a checkerboard):
35
+ - keep: photos, video thumbnails, illustrations, screenshots of colorful UI, colorful logos or icons that remain clear on a dark background, avatars, product shots on their own backgrounds.
36
+ - dim: large, very bright images (mostly white backgrounds, bright illustrations) that would glare on a dark page; they are shown slightly dimmed.
37
+ - invert: dark-on-light or dark-on-transparent line art, diagrams, charts, black text images, and mostly-monochrome logos whose few colored parts still look right after a hue-preserving inversion.
38
+ - light: monochrome dark glyphs, icons and wordmarks on transparent backgrounds that should simply become near-white.
39
+ - plate: logos or images with important dark parts on transparency whose colors must not be altered (multicolor brand marks, partner logos); they get a subtle light rounded backing.
40
+ Think about the site logo especially carefully: it must stay recognizable and clearly visible.
41
+
42
+ Answer with JSON only, following the schema. Use "#rrggbb", or "#rrggbbaa" when alpha is below 1."""
43
+
44
+ SCHEMA = {
45
+ "type": "object",
46
+ "properties": {
47
+ "design_notes": {"type": "string"},
48
+ "colors": {"type": "array", "items": {"type": "object", "properties": {"id": {"type": "string"}, "dark": {"type": "string"}},
49
+ "required": ["id", "dark"], "additionalProperties": False}},
50
+ "images": {"type": "array", "items": {"type": "object", "properties": {"id": {"type": "string"}, "treatment": {"type": "string", "enum": TREATMENTS}},
51
+ "required": ["id", "treatment"], "additionalProperties": False}},
52
+ },
53
+ "required": ["design_notes", "colors", "images"],
54
+ "additionalProperties": False,
55
+ }
56
+
57
+
58
+ def hexa(rgba):
59
+ r, g, b, a = rgba
60
+ h = "#%02x%02x%02x" % (r, g, b)
61
+ if a < 0.999:
62
+ h += "%02x" % round(a * 255)
63
+ return h
64
+
65
+
66
+ def pct(x):
67
+ if x >= 0.095:
68
+ return f"{x * 100:.0f}%"
69
+ if x >= 0.0005:
70
+ return f"{x * 100:.1f}%"
71
+ return "<0.1%"
72
+
73
+
74
+ def describe_tokens(page, max_tokens):
75
+ toks = page["tokens"][:max_tokens]
76
+ n = len(toks)
77
+ pa = max(1.0, page["docW"] * min(page["docH"], page["vh"] * 8))
78
+ va = page["vw"] * page["vh"]
79
+ # relations: (type, a, b, w)
80
+ best = {}
81
+ for typ, a, b, w in page["rel"]:
82
+ if a >= n or b >= n:
83
+ continue
84
+ best.setdefault((typ, a), []).append((w, b))
85
+ if typ in (0, 1):
86
+ best.setdefault((10 + typ, b), []).append((w, a)) # reverse: what sits on b
87
+ def top(key, k=1):
88
+ lst = sorted(best.get(key, []), reverse=True)[:k]
89
+ return [f"t{b}" for _, b in lst]
90
+ lines = []
91
+ for i, t in enumerate(toks):
92
+ d = t["d"]
93
+ role = ROLES[t["role"]]
94
+ parts = [f"t{i}", role, hexa(t["rgba"])]
95
+ info = []
96
+ if d["root"]:
97
+ info.append("PAGE CANVAS (html/body background)")
98
+ visible = d["nVis"] > 0
99
+ if role == "bg":
100
+ if d["area"]:
101
+ info.append(f"covers {pct(min(1, d['area'] / pa))} of page (unobstructed {pct(min(1, d['own'] / pa))}), first screen {pct(min(1, d['area1'] / va))}")
102
+ on = top((1, i))
103
+ if on and not d["root"]:
104
+ info.append(f"sits on {on[0]}")
105
+ txt = top((10, i), 3)
106
+ if txt:
107
+ info.append("text on it: " + ",".join(txt))
108
+ elif role in ("text", "fill", "stroke"):
109
+ if role == "text" and d["chars"]:
110
+ info.append(f"{d['chars']} chars, ~{d['fs']:.0f}px")
111
+ elif role != "text":
112
+ info.append("SVG " + ("fill (icons/logos/shapes)" if role == "fill" else "stroke (icons/lines)"))
113
+ on = top((0, i))
114
+ if on:
115
+ info.append(f"on {on[0]}" + (f" (contrast {d['contrast']:.1f}:1)" if d.get("contrast") else ""))
116
+ elif role == "border":
117
+ on = top((2, i))
118
+ if on:
119
+ info.append(f"line on {on[0]}")
120
+ elif role == "shadow":
121
+ info.append("box/text-shadow color")
122
+ elif role == "gradient":
123
+ info.append("gradient stop in a background-image")
124
+ if d["n"]:
125
+ info.append(f"{d['n']} elements" + ("" if visible else ", none visible now (hidden menus etc.)"))
126
+ flags = []
127
+ n_el = max(1, d["n"])
128
+ if d.get("media", 0) / n_el >= 0.5:
129
+ info.append("ON TOP OF / CONTAINS video or images (player UI, caption, scrim)")
130
+ for key, label in (("inter", "interactive"), ("link", "in links"), ("head", "headings"), ("header", "header"), ("nav", "nav"), ("footer", "footer")):
131
+ if d.get(key, 0) / n_el >= 0.3:
132
+ flags.append(label)
133
+ if flags:
134
+ info.append("/".join(flags))
135
+ if d["rules"] and not d["n"]:
136
+ info.append(f"CSS-rule only ({d['rules']} rules" + (f", {d['state']} hover/focus/active" if d["state"] else "") + ")")
137
+ elif d["state"]:
138
+ info.append(f"also in {d['state']} hover/focus/active rules")
139
+ if d.get("ex"):
140
+ info.append("e.g. " + "; ".join(x[:70] for x in d["ex"][:2]))
141
+ lines.append(" ".join(parts) + " | " + " | ".join(info))
142
+ return lines
143
+
144
+
145
+ def describe_images(page, max_images):
146
+ out = []
147
+ for j, im in enumerate(page["images"][:max_images]):
148
+ s = im.get("stats") or {}
149
+ bits = [f"i{j}", f"{im['w']}x{im['h']}px", im.get("ext") or im["kind"]]
150
+ if im["kind"] == "bg":
151
+ bits.append("CSS background image")
152
+ where = []
153
+ if im.get("header"): where.append("header/nav")
154
+ if im.get("footer"): where.append("footer")
155
+ if im.get("homeLink"): where.append("links to homepage")
156
+ elif im.get("inLink"): where.append("inside a link")
157
+ if im["y"] < page["vh"]: where.append("first screen")
158
+ if where: bits.append(", ".join(where))
159
+ if im.get("logo"): bits.append("named like a logo")
160
+ if im.get("icon"): bits.append("named like an icon")
161
+ if im.get("alt"): bits.append(f'alt "{im["alt"][:50]}"')
162
+ if s:
163
+ bits.append(f"transparent {pct(s.get('transp', 0))}, opaque mean lightness {s.get('meanL', 0):.2f}, near-white {pct(s.get('white', 0))}, near-black {pct(s.get('black', 0))}, chroma {s.get('meanC', 0):.2f}")
164
+ else:
165
+ bits.append("pixels unavailable")
166
+ if im.get("underTok", -1) >= 0:
167
+ bits.append(f"on surface t{im['underTok']}")
168
+ if im.get("count", 1) > 1:
169
+ bits.append(f"used {im['count']}x")
170
+ out.append(" | ".join(bits))
171
+ return out
172
+
173
+
174
+ def contact_sheet(page, max_images):
175
+ ims = [im for im in page["images"][:max_images]]
176
+ if not ims:
177
+ return None
178
+ cols = 6
179
+ cell, pad, lab = 112, 8, 16
180
+ rows = math.ceil(len(ims) / cols)
181
+ W, H = cols * (cell + pad) + pad, rows * (cell + pad + lab) + pad
182
+ sheet = Image.new("RGB", (W, H), (255, 255, 255))
183
+ dr = ImageDraw.Draw(sheet)
184
+ try:
185
+ font = ImageFont.load_default(size=13)
186
+ except Exception:
187
+ font = ImageFont.load_default()
188
+ for j, im in enumerate(ims):
189
+ x = pad + (j % cols) * (cell + pad)
190
+ y = pad + (j // cols) * (cell + pad + lab)
191
+ # checkerboard
192
+ for cy in range(0, cell, 8):
193
+ for cx in range(0, cell, 8):
194
+ c = (204, 204, 204) if ((cx + cy) // 8) % 2 else (245, 245, 245)
195
+ dr.rectangle([x + cx, y + lab + cy, x + cx + 7, y + lab + cy + 7], fill=c)
196
+ dr.text((x, y), f"i{j}", fill=(200, 0, 0), font=font)
197
+ if im.get("big"):
198
+ try:
199
+ b = base64.b64decode(im["big"].split(",", 1)[1])
200
+ th = Image.open(io.BytesIO(b)).convert("RGBA").resize((cell, cell), Image.LANCZOS)
201
+ sheet.paste(th, (x, y + lab), th)
202
+ except Exception:
203
+ dr.text((x + 4, y + lab + 40), "n/a", fill=(0, 0, 0), font=font)
204
+ else:
205
+ dr.text((x + 4, y + lab + 40), "n/a", fill=(0, 0, 0), font=font)
206
+ buf = io.BytesIO()
207
+ sheet.save(buf, "PNG", optimize=True)
208
+ return buf.getvalue()
209
+
210
+
211
+ def screenshot(pdir, max_w=900, max_h=1400):
212
+ im = Image.open(pdir / "shot.jpg").convert("RGB")
213
+ s = max_w / im.width
214
+ im = im.resize((max_w, int(im.height * s)), Image.LANCZOS)
215
+ if im.height > max_h:
216
+ im = im.crop((0, 0, max_w, max_h))
217
+ buf = io.BytesIO()
218
+ im.save(buf, "JPEG", quality=82)
219
+ return buf.getvalue()
220
+
221
+
222
+ def build_request(page, pdir, args):
223
+ lines = describe_tokens(page, args.max_tokens)
224
+ imgs = describe_images(page, args.max_images)
225
+ sheet = contact_sheet(page, args.max_images) if imgs else None
226
+ host = re.sub(r"^https?://", "", page["url"]).split("/")[0]
227
+ text = [f"Website: {host} — \"{page['title'][:100]}\"", f"Viewport 1280x{page['vh']}, page height {page['docH']}px.",
228
+ "", f"COLOR INVENTORY ({len(lines)} colors; id role color | usage):"] + lines
229
+ if imgs:
230
+ text += ["", f"IMAGES ({len(imgs)}; thumbnails in the contact sheet):"] + imgs
231
+ text += ["", f"Return the dark-theme color for every id t0-t{len(lines) - 1}" + (f" and a treatment for every image i0-i{len(imgs) - 1}." if imgs else " (images: return an empty list).")]
232
+ content = [{"type": "text", "text": "Screenshot of the light page (top section):"},
233
+ {"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": base64.b64encode(screenshot(pdir)).decode()}}]
234
+ if sheet:
235
+ content += [{"type": "text", "text": "Contact sheet of the page's images:"},
236
+ {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": base64.b64encode(sheet).decode()}}]
237
+ content.append({"type": "text", "text": "\n".join(text)})
238
+ return content, len(lines), len(imgs)
239
+
240
+
241
+ HEX_RE = re.compile(r"^#?([0-9a-fA-F]{6})([0-9a-fA-F]{2})?$")
242
+
243
+
244
+ def parse_answer(msg, n_tok, n_img):
245
+ txt = "".join(b.text for b in msg.content if b.type == "text")
246
+ data = json.loads(txt)
247
+ colors, images = {}, {}
248
+ for c in data.get("colors", []):
249
+ m = HEX_RE.match(c.get("dark", "").strip())
250
+ mid = re.match(r"^t(\d+)$", c.get("id", ""))
251
+ if m and mid and int(mid.group(1)) < n_tok:
252
+ colors[int(mid.group(1))] = "#" + m.group(1).lower() + (m.group(2) or "").lower()
253
+ for c in data.get("images", []):
254
+ mid = re.match(r"^i(\d+)$", c.get("id", ""))
255
+ if mid and int(mid.group(1)) < n_img and c.get("treatment") in TREATMENTS:
256
+ images[int(mid.group(1))] = c["treatment"]
257
+ return data.get("design_notes", ""), colors, images
258
+
259
+
260
+ class Gate:
261
+ """Adaptive concurrency: grow slowly while healthy, back off on the org concurrency cap and stay below it."""
262
+ def __init__(self, start, cap):
263
+ self.limit, self.cap, self.active = float(start), cap, 0
264
+ self.cond = asyncio.Condition()
265
+ self.last429 = 0.0
266
+
267
+ async def acquire(self):
268
+ async with self.cond:
269
+ while self.active >= int(self.limit):
270
+ await self.cond.wait()
271
+ self.active += 1
272
+
273
+ async def release(self, ok=True, limited=False):
274
+ async with self.cond:
275
+ self.active -= 1
276
+ if limited:
277
+ seen = self.active + 1
278
+ self.cap = max(2, min(self.cap, int(seen * 0.75)))
279
+ self.limit = max(2.0, min(self.limit * 0.7, self.cap))
280
+ self.last429 = time.time()
281
+ elif ok and time.time() - self.last429 > 120:
282
+ self.limit = min(self.cap, self.limit + 0.2)
283
+ self.cond.notify_all()
284
+
285
+
286
+ class Ledger:
287
+ def __init__(self, path, budget):
288
+ self.path, self.budget = Path(path), budget
289
+ self.cost = 0.0
290
+ self.n = 0
291
+ if self.path.exists():
292
+ for line in self.path.open():
293
+ try:
294
+ r = json.loads(line)
295
+ self.cost += r["cost"]; self.n += 1
296
+ except Exception:
297
+ pass
298
+ self.f = self.path.open("a")
299
+
300
+ def add(self, rec):
301
+ self.cost += rec["cost"]; self.n += 1
302
+ self.f.write(json.dumps(rec) + "\n"); self.f.flush()
303
+
304
+
305
+ def usage_cost(u):
306
+ cr = getattr(u, "cache_read_input_tokens", 0) or 0
307
+ cw = getattr(u, "cache_creation_input_tokens", 0) or 0
308
+ return u.input_tokens * PRICE_IN + u.output_tokens * PRICE_OUT + cr * PRICE_CACHE_READ + cw * PRICE_CACHE_WRITE
309
+
310
+
311
+ ERRLOG = {}
312
+
313
+
314
+ def logerr(kind, e):
315
+ ERRLOG[kind] = ERRLOG.get(kind, 0) + 1
316
+ if ERRLOG[kind] <= 3:
317
+ print(f"[{kind}] {type(e).__name__}: {str(e)[:300]}", flush=True)
318
+
319
+
320
+ async def one(client, sem, pid, crawl_dir, out_dir, ledger, args, stats):
321
+ out = out_dir / f"{pid}.json"
322
+ if out.exists():
323
+ return
324
+ pdir = crawl_dir / "pages" / pid
325
+ try:
326
+ page = json.load(gzip.open(pdir / "light.json.gz"))
327
+ content, n_tok, n_img = build_request(page, pdir, args)
328
+ except Exception as e:
329
+ stats["prep_err"] = stats.get("prep_err", 0) + 1
330
+ return
331
+ for attempt in range(8):
332
+ if ledger.cost >= args.budget:
333
+ return
334
+ await sem.acquire()
335
+ ok, limited = False, False
336
+ t0 = time.time()
337
+ try:
338
+ async with client.messages.stream(
339
+ model=MODEL, max_tokens=args.max_out, system=SYSTEM,
340
+ thinking={"type": "adaptive"}, output_config={"effort": args.effort, "format": {"type": "json_schema", "schema": SCHEMA}},
341
+ messages=[{"role": "user", "content": content}],
342
+ ) as stream:
343
+ msg = await stream.get_final_message()
344
+ ok = True
345
+ except anthropic.RateLimitError as e:
346
+ stats["429"] = stats.get("429", 0) + 1
347
+ limited = "concurrent" in str(e).lower()
348
+ logerr("429", e)
349
+ except anthropic.APIStatusError as e:
350
+ stats[f"http{e.status_code}"] = stats.get(f"http{e.status_code}", 0) + 1
351
+ logerr(f"http{e.status_code}", e)
352
+ if e.status_code in (400, 401, 403, 404, 413):
353
+ (out_dir / "errors").mkdir(exist_ok=True)
354
+ (out_dir / "errors" / f"{pid}.txt").write_text(str(e)[:2000])
355
+ await sem.release(False)
356
+ return
357
+ except (anthropic.APIConnectionError, anthropic.APITimeoutError) as e:
358
+ stats["conn"] = stats.get("conn", 0) + 1
359
+ logerr("conn", e)
360
+ except Exception as e:
361
+ stats["other"] = stats.get("other", 0) + 1
362
+ logerr("other", e)
363
+ await sem.release(ok, limited)
364
+ if not ok:
365
+ await asyncio.sleep(min(60, 4 * (attempt + 1)) + random.random() * 4)
366
+ continue
367
+ cost = usage_cost(msg.usage)
368
+ rec = {"id": pid, "cost": cost, "in": msg.usage.input_tokens, "out": msg.usage.output_tokens, "stop": msg.stop_reason,
369
+ "secs": round(time.time() - t0, 1), "t": time.time()}
370
+ ledger.add(rec)
371
+ if msg.stop_reason == "refusal":
372
+ stats["refusal"] = stats.get("refusal", 0) + 1
373
+ out.write_text(json.dumps({"id": pid, "refusal": True, "usage": rec}))
374
+ return
375
+ try:
376
+ notes, colors, images = parse_answer(msg, n_tok, n_img)
377
+ except Exception as e:
378
+ stats["parse"] = stats.get("parse", 0) + 1
379
+ (out_dir / "errors").mkdir(exist_ok=True)
380
+ (out_dir / "errors" / f"{pid}.txt").write_text("parse: " + str(e)[:500])
381
+ return
382
+ out.write_text(json.dumps({"id": pid, "url": page["url"], "n_tok": n_tok, "n_img": n_img, "notes": notes,
383
+ "colors": colors, "images": images, "usage": rec, "effort": args.effort, "model": MODEL}))
384
+ stats["ok"] = stats.get("ok", 0) + 1
385
+ return
386
+
387
+
388
+ # ---------------- batch mode ----------------
389
+ import urllib.request
390
+
391
+ PROXY = os.environ.get("ANTHROPIC_BASE_URL", "https://api.anthropic.com")
392
+ HDRS = {"x-api-key": os.environ.get("ANTHROPIC_API_KEY", ""), "anthropic-version": "2023-06-01", "content-type": "application/json"}
393
+
394
+
395
+ def http(method, path, body=None, timeout=300):
396
+ req = urllib.request.Request(PROXY + path, data=json.dumps(body).encode() if body is not None else None, headers=HDRS, method=method)
397
+ with urllib.request.urlopen(req, timeout=timeout) as r:
398
+ return r.read()
399
+
400
+
401
+ def params_for(content, args):
402
+ return {"model": MODEL, "max_tokens": args.max_out, "system": SYSTEM, "thinking": {"type": "adaptive"},
403
+ "output_config": {"effort": args.effort, "format": {"type": "json_schema", "schema": SCHEMA}},
404
+ "messages": [{"role": "user", "content": content}]}
405
+
406
+
407
+ class Msg: # minimal adapter so parse_answer works on batch result dicts
408
+ def __init__(self, d):
409
+ self.content = [type("B", (), {"type": b.get("type"), "text": b.get("text", "")}) for b in d.get("content", [])]
410
+
411
+
412
+ def collect_batch(bid, meta, crawl_dir, out_dir, ledger, stats):
413
+ raw = http("GET", f"/v1/messages/batches/{bid}/results", timeout=600).decode()
414
+ seen = set()
415
+ for line in raw.splitlines():
416
+ if not line.strip():
417
+ continue
418
+ r = json.loads(line)
419
+ pid = r["custom_id"]
420
+ seen.add(pid)
421
+ res = r["result"]
422
+ if res["type"] != "succeeded":
423
+ stats["b_" + res["type"]] = stats.get("b_" + res["type"], 0) + 1
424
+ continue
425
+ m = res["message"]
426
+ u = m.get("usage", {})
427
+ cost = 0.5 * (u.get("input_tokens", 0) * PRICE_IN + u.get("output_tokens", 0) * PRICE_OUT
428
+ + (u.get("cache_read_input_tokens") or 0) * PRICE_CACHE_READ + (u.get("cache_creation_input_tokens") or 0) * PRICE_CACHE_WRITE)
429
+ rec = {"id": pid, "cost": cost, "in": u.get("input_tokens", 0), "out": u.get("output_tokens", 0), "stop": m.get("stop_reason"), "batch": bid, "t": time.time()}
430
+ ledger.add(rec)
431
+ out = out_dir / f"{pid}.json"
432
+ if m.get("stop_reason") == "refusal":
433
+ stats["refusal"] = stats.get("refusal", 0) + 1
434
+ out.write_text(json.dumps({"id": pid, "refusal": True, "usage": rec}))
435
+ continue
436
+ n_tok, n_img = meta[pid]
437
+ try:
438
+ notes, colors, images = parse_answer(Msg(m), n_tok, n_img)
439
+ except Exception as e:
440
+ stats["parse"] = stats.get("parse", 0) + 1
441
+ continue
442
+ url = ""
443
+ out.write_text(json.dumps({"id": pid, "n_tok": n_tok, "n_img": n_img, "notes": notes, "colors": colors, "images": images,
444
+ "usage": rec, "effort": args_global.effort, "model": MODEL}))
445
+ stats["ok"] = stats.get("ok", 0) + 1
446
+ return seen
447
+
448
+
449
+ args_global = None
450
+
451
+
452
+ def batch_main(args):
453
+ global args_global
454
+ args_global = args
455
+ crawl_dir, out_dir = Path(args.crawl), Path(args.out)
456
+ out_dir.mkdir(parents=True, exist_ok=True)
457
+ ledger = Ledger(out_dir / "ledger.jsonl", args.budget)
458
+ state_p = out_dir / "batches.jsonl"
459
+ batches = {}
460
+ if state_p.exists():
461
+ for line in state_p.open():
462
+ b = json.loads(line)
463
+ batches.setdefault(b["id"], {}).update(b)
464
+ submitted = set(pid for b in batches.values() for pid in b.get("ids", []))
465
+ stats = {}
466
+ est_per = 0.0035
467
+ last_submit = 0
468
+ while True:
469
+ # collect finished batches
470
+ pending_est = 0.0
471
+ for bid, b in list(batches.items()):
472
+ if b.get("collected"):
473
+ continue
474
+ try:
475
+ info = json.loads(http("GET", f"/v1/messages/batches/{bid}"))
476
+ except Exception as e:
477
+ print("poll error", bid, str(e)[:200], flush=True)
478
+ pending_est += len(b["ids"]) * est_per
479
+ continue
480
+ if info["processing_status"] == "ended":
481
+ try:
482
+ collect_batch(bid, b["meta"], crawl_dir, out_dir, ledger, stats)
483
+ b["collected"] = True
484
+ with state_p.open("a") as f:
485
+ f.write(json.dumps({"id": bid, "collected": True}) + "\n")
486
+ print(f"[{time.strftime('%H:%M:%S')}] collected {bid} {info['request_counts']}", flush=True)
487
+ except Exception as e:
488
+ print("collect error", bid, str(e)[:300], flush=True)
489
+ pending_est += len(b["ids"]) * est_per
490
+ else:
491
+ pending_est += len(b["ids"]) * est_per
492
+ if ledger.n > 50:
493
+ est_per = max(0.002, ledger.cost / ledger.n)
494
+ # candidates
495
+ cands = []
496
+ if args.ids:
497
+ cands = [l.strip() for l in open(args.ids) if l.strip()]
498
+ else:
499
+ for line in (l for ip in sorted(crawl_dir.glob("index*.jsonl")) for l in open(ip)):
500
+ try:
501
+ r = json.loads(line)
502
+ except Exception:
503
+ continue
504
+ if r.get("kind") in args.kinds.split(",") and (r.get("n_tokens") or 0) >= 8 and (r.get("status") or 200) < 400:
505
+ cands.append(r["id"])
506
+ new = [c for c in cands if c not in submitted and not (out_dir / f"{c}.json").exists()]
507
+ room = args.budget - ledger.cost - pending_est
508
+ can = int(max(0, room) / est_per)
509
+ ready = len(new) >= args.batch_size or (new and time.time() - last_submit > args.max_wait)
510
+ if ready and can > 0:
511
+ chunk = new[: min(args.batch_size, can)]
512
+ reqs, meta = [], {}
513
+ for pid in chunk:
514
+ pdir = crawl_dir / "pages" / pid
515
+ try:
516
+ page = json.load(gzip.open(pdir / "light.json.gz"))
517
+ content, n_tok, n_img = build_request(page, pdir, args)
518
+ except Exception as e:
519
+ submitted.add(pid)
520
+ continue
521
+ reqs.append({"custom_id": pid, "params": params_for(content, args)})
522
+ meta[pid] = (n_tok, n_img)
523
+ if reqs:
524
+ try:
525
+ info = json.loads(http("POST", "/v1/messages/batches", {"requests": reqs}, timeout=900))
526
+ b = {"id": info["id"], "ids": list(meta), "meta": meta, "t": time.time()}
527
+ batches[info["id"]] = b
528
+ with state_p.open("a") as f:
529
+ f.write(json.dumps(b) + "\n")
530
+ submitted.update(meta)
531
+ last_submit = time.time()
532
+ print(f"[{time.strftime('%H:%M:%S')}] submitted {info['id']} with {len(reqs)} (spent ${ledger.cost:.2f}, pending ~${pending_est:.2f}, est/page ${est_per:.4f})", flush=True)
533
+ except Exception as e:
534
+ print("submit error", str(e)[:300], flush=True)
535
+ time.sleep(60)
536
+ print(f"[{time.strftime('%H:%M:%S')}] spent ${ledger.cost:.3f} on {ledger.n} | open batches {sum(1 for b in batches.values() if not b.get('collected'))} | waiting pages {len(new)} | {stats}", flush=True)
537
+ open_batches = [b for b in batches.values() if not b.get("collected")]
538
+ if not args.follow and not open_batches and (not new or can <= 0):
539
+ break
540
+ if room <= 0 and not open_batches:
541
+ print("budget reached", flush=True)
542
+ break
543
+ time.sleep(args.poll)
544
+ print(f"DONE ${ledger.cost:.3f} {ledger.n} {stats}", flush=True)
545
+
546
+
547
+ async def main():
548
+ ap = argparse.ArgumentParser()
549
+ ap.add_argument("--crawl", required=True)
550
+ ap.add_argument("--out", required=True)
551
+ ap.add_argument("--ids", help="file with page ids (default: all light pages in index)")
552
+ ap.add_argument("--budget", type=float, default=95.0)
553
+ ap.add_argument("--concurrency", type=int, default=64)
554
+ ap.add_argument("--effort", default="high")
555
+ ap.add_argument("--max-tokens", type=int, default=140)
556
+ ap.add_argument("--max-images", type=int, default=24)
557
+ ap.add_argument("--max-out", type=int, default=48000)
558
+ ap.add_argument("--limit", type=int, default=0)
559
+ ap.add_argument("--follow", action="store_true", help="keep polling the crawl index for new pages")
560
+ ap.add_argument("--mode", default="batch", choices=["batch", "realtime"])
561
+ ap.add_argument("--batch-size", type=int, default=250)
562
+ ap.add_argument("--max-wait", type=int, default=900)
563
+ ap.add_argument("--poll", type=int, default=60)
564
+ ap.add_argument("--kinds", default="light")
565
+ args = ap.parse_args()
566
+ if args.mode == "batch":
567
+ return batch_main(args)
568
+ crawl_dir, out_dir = Path(args.crawl), Path(args.out)
569
+ out_dir.mkdir(parents=True, exist_ok=True)
570
+ ledger = Ledger(out_dir / "ledger.jsonl", args.budget)
571
+ client = anthropic.AsyncAnthropic(base_url=PROXY, api_key=os.environ.get("ANTHROPIC_API_KEY", ""),
572
+ max_retries=0, timeout=600)
573
+ sem = Gate(min(8, args.concurrency), args.concurrency)
574
+ stats = {}
575
+ started = set()
576
+ tasks = []
577
+ t_start = time.time()
578
+
579
+ def candidates():
580
+ if args.ids:
581
+ return [l.strip() for l in open(args.ids) if l.strip()]
582
+ ids = []
583
+ for line in (l for ip in sorted(crawl_dir.glob("index*.jsonl")) for l in open(ip)):
584
+ try:
585
+ r = json.loads(line)
586
+ except Exception:
587
+ continue
588
+ if r.get("kind") == "light" and (r.get("n_tokens") or 0) >= 8 and (r.get("status") or 200) < 400:
589
+ ids.append(r["id"])
590
+ return ids
591
+
592
+ async def reporter():
593
+ while True:
594
+ await asyncio.sleep(60)
595
+ el = time.time() - t_start
596
+ print(f"[{time.strftime('%H:%M:%S')}] ${ledger.cost:.3f} spent, {ledger.n} calls, {stats}, {ledger.n / max(1, el) * 3600:.0f}/h, conc {sem.active}/{sem.limit:.1f} cap {sem.cap}", flush=True)
597
+ rep = asyncio.create_task(reporter())
598
+ while True:
599
+ ids = candidates()
600
+ new = [i for i in ids if i not in started]
601
+ if args.limit:
602
+ new = new[: max(0, args.limit - len(started))]
603
+ for pid in new:
604
+ started.add(pid)
605
+ tasks.append(asyncio.create_task(one(client, sem, pid, crawl_dir, out_dir, ledger, args, stats)))
606
+ tasks = [t for t in tasks if not t.done()]
607
+ if ledger.cost >= args.budget:
608
+ print("budget reached", flush=True)
609
+ break
610
+ if not args.follow and not tasks:
611
+ break
612
+ if args.limit and len(started) >= args.limit and not tasks:
613
+ break
614
+ await asyncio.sleep(30 if args.follow else 2)
615
+ await asyncio.gather(*tasks, return_exceptions=True)
616
+ rep.cancel()
617
+ print(f"DONE ${ledger.cost:.3f} {ledger.n} calls {stats}", flush=True)
618
+
619
+
620
+ if __name__ == "__main__":
621
+ asyncio.run(main())
training/judge.py ADDED
@@ -0,0 +1,110 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """Blind pairwise judging of dark-mode renders with Claude Haiku 5.5 (Batch API). Runs on Vega's host."""
3
+ import argparse, base64, io, json, random, time
4
+ from pathlib import Path
5
+ from PIL import Image
6
+
7
+ import haiku_gen as hg
8
+
9
+ JUDGE_SYSTEM = """You are a meticulous product-design reviewer. You compare two dark-mode versions (A and B) of the same web page, given the original light page for reference. Judge which dark mode a demanding user would rather browse with, considering in this order: (1) readability: every piece of text clearly legible against its background, nothing lost; (2) nothing broken: no large light patches left glaring, logos and icons still visible and recognizable, photos and videos not color-distorted; (3) design quality: surfaces, hierarchy and elevation preserved, brand colors recognizable, pleasant and coherent palette, no muddy or harsh colors. Small differences in exact shades matter less than legibility and breakage. Answer with JSON only."""
10
+
11
+ SCHEMA = {"type": "object", "properties": {
12
+ "winner": {"type": "string", "enum": ["A", "B", "tie"]},
13
+ "score_A": {"type": "integer"}, "score_B": {"type": "integer"},
14
+ "reason": {"type": "string"}},
15
+ "required": ["winner", "score_A", "score_B", "reason"], "additionalProperties": False}
16
+
17
+
18
+ def img_block(path, w=800, h=1000):
19
+ im = Image.open(path).convert("RGB")
20
+ s = w / im.width
21
+ im = im.resize((w, int(im.height * s)), Image.LANCZOS).crop((0, 0, w, min(h, int(im.height * s))))
22
+ buf = io.BytesIO(); im.save(buf, "JPEG", quality=80)
23
+ return {"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": base64.b64encode(buf.getvalue()).decode()}}
24
+
25
+
26
+ def main():
27
+ ap = argparse.ArgumentParser()
28
+ ap.add_argument("--renders", required=True, help="render output dir with one subdir per tag")
29
+ ap.add_argument("--crawl", required=True)
30
+ ap.add_argument("--pairs", default="nadir:dr,nadir:haiku")
31
+ ap.add_argument("--out", required=True)
32
+ ap.add_argument("--effort", default="medium")
33
+ ap.add_argument("--budget", type=float, default=3.0)
34
+ ap.add_argument("--chunk", type=int, default=400, help="requests per batch (3 images each, ~150 KB)")
35
+ args = ap.parse_args()
36
+ R, crawl, out = Path(args.renders), Path(args.crawl), Path(args.out)
37
+ out.mkdir(parents=True, exist_ok=True)
38
+ rng = random.Random(0)
39
+ reqs, meta = [], {}
40
+ for pair in args.pairs.split(","):
41
+ x, y = pair.split(":")
42
+ for rx in sorted((R / x).glob("*.json")):
43
+ pid = rx.stem
44
+ ry = R / y / f"{pid}.json"
45
+ if not ry.exists():
46
+ continue
47
+ jx, jy = json.loads(rx.read_text()), json.loads(ry.read_text())
48
+ if not (jx.get("ok") and jy.get("ok")) or not (R / x / f"{pid}.jpg").exists() or not (R / y / f"{pid}.jpg").exists():
49
+ continue
50
+ flip = rng.random() < 0.5
51
+ a, b = (y, x) if flip else (x, y)
52
+ cid = f"{pid}-{x}-{y}"
53
+ content = [{"type": "text", "text": "Original light page:"}, img_block(crawl / "pages" / pid / "shot.jpg"),
54
+ {"type": "text", "text": "Dark mode A:"}, img_block(R / a / f"{pid}.jpg"),
55
+ {"type": "text", "text": "Dark mode B:"}, img_block(R / b / f"{pid}.jpg"),
56
+ {"type": "text", "text": "Which dark mode is better? Scores are 1-10."}]
57
+ reqs.append({"custom_id": cid, "params": {"model": hg.MODEL, "max_tokens": 16000, "system": JUDGE_SYSTEM, "thinking": {"type": "adaptive"},
58
+ "output_config": {"effort": args.effort, "format": {"type": "json_schema", "schema": SCHEMA}},
59
+ "messages": [{"role": "user", "content": content}]}})
60
+ meta[cid] = {"pid": pid, "x": x, "y": y, "A": a, "B": b}
61
+ print(len(reqs), "comparisons", flush=True)
62
+ est = len(reqs) * 0.0015
63
+ if est > args.budget:
64
+ reqs = reqs[: int(args.budget / 0.0015)]
65
+ results, cost = [], 0.0
66
+ bids = []
67
+ for k in range(0, len(reqs), args.chunk): # submit every chunk first, then wait for all of them together
68
+ info = json.loads(hg.http("POST", "/v1/messages/batches", {"requests": reqs[k:k + args.chunk]}, timeout=900))
69
+ bids.append(info["id"])
70
+ print("submitted", info["id"], len(reqs[k:k + args.chunk]), flush=True)
71
+ pending = set(bids)
72
+ while pending:
73
+ time.sleep(30)
74
+ for bid in list(pending):
75
+ if json.loads(hg.http("GET", f"/v1/messages/batches/{bid}"))["processing_status"] == "ended":
76
+ pending.discard(bid)
77
+ for bid in bids:
78
+ raw = hg.http("GET", f"/v1/messages/batches/{bid}/results", timeout=600).decode()
79
+ for line in raw.splitlines():
80
+ if not line.strip():
81
+ continue
82
+ r = json.loads(line)
83
+ if r["result"]["type"] != "succeeded":
84
+ continue
85
+ m = r["result"]["message"]; u = m["usage"]
86
+ cost += 0.5 * (u["input_tokens"] * hg.PRICE_IN + u["output_tokens"] * hg.PRICE_OUT)
87
+ txt = "".join(b.get("text", "") for b in m["content"] if b.get("type") == "text")
88
+ try:
89
+ v = json.loads(txt)
90
+ except Exception:
91
+ continue
92
+ mt = meta[r["custom_id"]]
93
+ win = {"A": mt["A"], "B": mt["B"], "tie": "tie"}[v["winner"]]
94
+ sc = {mt["A"]: v["score_A"], mt["B"]: v["score_B"]}
95
+ results.append({**mt, "winner": win, "score_x": sc[mt["x"]], "score_y": sc[mt["y"]], "reason": v["reason"]})
96
+ summary = {}
97
+ for pair in args.pairs.split(","):
98
+ x, y = pair.split(":")
99
+ rs = [r for r in results if r["x"] == x and r["y"] == y]
100
+ if not rs:
101
+ continue
102
+ wx = sum(r["winner"] == x for r in rs); wy = sum(r["winner"] == y for r in rs); t = sum(r["winner"] == "tie" for r in rs)
103
+ summary[pair] = {"n": len(rs), f"{x}_wins": wx, f"{y}_wins": wy, "ties": t, f"{x}_winrate_excl_ties": round(wx / max(1, wx + wy), 3),
104
+ f"mean_score_{x}": round(sum(r["score_x"] for r in rs) / len(rs), 2), f"mean_score_{y}": round(sum(r["score_y"] for r in rs) / len(rs), 2)}
105
+ json.dump({"summary": summary, "cost": cost, "results": results}, open(out / "judge.json", "w"), indent=1)
106
+ print(json.dumps(summary, indent=1), f"cost ${cost:.3f}", flush=True)
107
+
108
+
109
+ if __name__ == "__main__":
110
+ main()
training/model.py ADDED
@@ -0,0 +1,141 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """NadirNet: a relation-biased set transformer over a page's colour palette and images.
2
+
3
+ Tokens: [PAGE] + one token per (colour, role) the page uses + one token per image.
4
+ Attention logits get a learned per-head bias from the page's structure graph (text-on-surface,
5
+ surface-on-surface, border-on-surface, same-element, image-on-surface, same-colour), so the model
6
+ reasons about hierarchy and contrast the way a designer would. Outputs a dark colour per token
7
+ (OKLab, parameterised relative to the input so hue is kept unless the model chooses otherwise),
8
+ an alpha, and a treatment class per image.
9
+ """
10
+ import math
11
+ import torch
12
+ import torch.nn as nn
13
+ import torch.nn.functional as F
14
+
15
+ FT, FI, FP = 49, 39, 8
16
+ N_ROLES, N_TREAT = 7, 5
17
+ K_REL = 11
18
+
19
+
20
+ def fourier_lab(lab, n=6):
21
+ x = torch.stack([lab[..., 0], lab[..., 1] * 4, lab[..., 2] * 4], -1) # (..., 3)
22
+ freqs = (2.0 ** torch.arange(n, device=lab.device, dtype=lab.dtype)) * math.pi
23
+ ang = x[..., None] * freqs # (...,3,n)
24
+ return torch.cat([torch.sin(ang), torch.cos(ang)], -1).flatten(-2) # (..., 6n)
25
+
26
+
27
+ class MLP(nn.Module):
28
+ def __init__(self, i, h, o, drop=0.0):
29
+ super().__init__()
30
+ self.net = nn.Sequential(nn.Linear(i, h), nn.GELU(), nn.Dropout(drop), nn.Linear(h, o))
31
+
32
+ def forward(self, x):
33
+ return self.net(x)
34
+
35
+
36
+ class ThumbCNN(nn.Module):
37
+ """32x32 RGBA thumbnail -> embedding. Alpha matters: transparent logos behave differently from photos."""
38
+ def __init__(self, o=128):
39
+ super().__init__()
40
+ c = [5, 32, 64, 96, 128]
41
+ layers = []
42
+ for a, b in zip(c[:-1], c[1:]):
43
+ layers += [nn.Conv2d(a, b, 3, padding=1), nn.GELU(), nn.Conv2d(b, b, 3, padding=1, stride=2), nn.GELU()]
44
+ self.net = nn.Sequential(*layers)
45
+ self.proj = nn.Linear(128 * 2 * 2, o)
46
+
47
+ def forward(self, x): # (B, M, 4, 32, 32) in [0,1]
48
+ B, M = x.shape[:2]
49
+ x = x.reshape(B * M, 4, 32, 32)
50
+ a = x[:, 3:4]
51
+ lum = (0.2126 * x[:, 0:1] + 0.7152 * x[:, 1:2] + 0.0722 * x[:, 2:3]) * a # premultiplied luminance
52
+ x = torch.cat([x[:, :3] * a, a, lum], 1)
53
+ x = self.net(x) # (BM,128,2,2)
54
+ return self.proj(x.flatten(1)).reshape(B, M, -1)
55
+
56
+
57
+ class Block(nn.Module):
58
+ def __init__(self, d, h, ff, drop):
59
+ super().__init__()
60
+ self.h = h
61
+ self.ln1 = nn.LayerNorm(d)
62
+ self.qkv = nn.Linear(d, 3 * d)
63
+ self.proj = nn.Linear(d, d)
64
+ self.rel = nn.Linear(K_REL, h, bias=False)
65
+ nn.init.normal_(self.rel.weight, std=0.5)
66
+ self.ln2 = nn.LayerNorm(d)
67
+ self.ff = nn.Sequential(nn.Linear(d, ff * d), nn.GELU(), nn.Dropout(drop), nn.Linear(ff * d, d))
68
+ self.drop = nn.Dropout(drop)
69
+
70
+ def forward(self, x, rel, pad_bias):
71
+ B, T, D = x.shape
72
+ q, k, v = self.qkv(self.ln1(x)).reshape(B, T, 3, self.h, D // self.h).permute(2, 0, 3, 1, 4)
73
+ bias = self.rel(rel).permute(0, 3, 1, 2) + pad_bias # (B,H,T,T)
74
+ a = F.scaled_dot_product_attention(q, k, v, attn_mask=bias.to(q.dtype))
75
+ x = x + self.drop(self.proj(a.transpose(1, 2).reshape(B, T, D)))
76
+ return x + self.drop(self.ff(self.ln2(x)))
77
+
78
+
79
+ class NadirNet(nn.Module):
80
+ def __init__(self, d=384, layers=8, heads=6, ff=4, drop=0.1):
81
+ super().__init__()
82
+ self.cfg = dict(d=d, layers=layers, heads=heads, ff=ff, drop=drop)
83
+ self.tok_in = MLP(FT + N_ROLES + 36, 2 * d, d)
84
+ self.cnn = ThumbCNN(128)
85
+ self.img_in = MLP(FI + 128, 2 * d, d)
86
+ self.page_in = MLP(FP, d, d)
87
+ self.src = nn.Embedding(2, d)
88
+ self.typ = nn.Embedding(3, d)
89
+ self.blocks = nn.ModuleList([Block(d, heads, ff, drop) for _ in range(layers)])
90
+ self.ln = nn.LayerNorm(d)
91
+ self.color_head = MLP(d + 3, d, 6)
92
+ self.img_head = MLP(d, d // 2, N_TREAT)
93
+
94
+ def forward(self, tok_f, tok_role, tok_lab, img_f, img_thumb, page_f, src, rel, mask):
95
+ """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,
96
+ page_f (B,FP), src (B,) long, rel (B,T,T,K) with T=1+N+M, mask (B,T) bool (True=real)."""
97
+ B, N = tok_role.shape
98
+ M = img_f.shape[1]
99
+ 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]
100
+ if M > 0:
101
+ img = self.img_in(torch.cat([img_f, self.cnn(img_thumb)], -1)) + self.typ.weight[2]
102
+ else:
103
+ img = tok.new_zeros(B, 0, tok.shape[-1])
104
+ cls = self.page_in(page_f)[:, None] + self.typ.weight[0]
105
+ x = torch.cat([cls, tok, img], 1) + self.src(src)[:, None]
106
+ pad_bias = torch.zeros(mask.shape, dtype=x.dtype, device=x.device).masked_fill(~mask, -1e4)[:, None, None, :]
107
+ for blk in self.blocks:
108
+ x = blk(x, rel, pad_bias)
109
+ x = self.ln(x)
110
+ ht = x[:, 1:1 + N]
111
+ z = self.color_head(torch.cat([ht, tok_lab], -1))
112
+ L = torch.sigmoid(z[..., 0])
113
+ s = 1.6 * torch.sigmoid(z[..., 1])
114
+ r = 0.6 * torch.tanh(z[..., 2])
115
+ a_in, b_in = tok_lab[..., 1], tok_lab[..., 2]
116
+ a = s * a_in - r * b_in + 0.12 * torch.tanh(z[..., 3])
117
+ b = r * a_in + s * b_in + 0.12 * torch.tanh(z[..., 4])
118
+ lab = torch.stack([L, a, b], -1)
119
+ alpha_logit = z[..., 5]
120
+ img_logits = self.img_head(x[:, 1 + N:]) if M > 0 else x.new_zeros(B, 0, N_TREAT)
121
+ return lab, alpha_logit, img_logits
122
+
123
+
124
+ # ---------- colour maths (differentiable) ----------
125
+ OKLAB_TO_LMS = torch.tensor([[1.0, 0.3963377774, 0.2158037573], [1.0, -0.1055613458, -0.0638541728], [1.0, -0.0894841775, -1.2914855480]])
126
+ LMS_TO_RGB = torch.tensor([[4.0767416621, -3.3077115913, 0.2309699292], [-1.2684380046, 2.6097574011, -0.3413193965], [-0.0041960863, -0.7034186147, 1.7076147010]])
127
+
128
+
129
+ def oklab_to_linrgb(lab):
130
+ lms = lab @ OKLAB_TO_LMS.to(lab).T
131
+ return (lms ** 3) @ LMS_TO_RGB.to(lab).T
132
+
133
+
134
+ def luminance(lab):
135
+ rgb = oklab_to_linrgb(lab).clamp(0, 1)
136
+ return rgb @ torch.tensor([0.2126, 0.7152, 0.0722]).to(lab)
137
+
138
+
139
+ def contrast_ratio(lab1, lab2):
140
+ l1, l2 = luminance(lab1), luminance(lab2)
141
+ return (torch.maximum(l1, l2) + 0.05) / (torch.minimum(l1, l2) + 0.05)
training/predict.py ADDED
@@ -0,0 +1,93 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """Run a trained NadirNet over held-out pages and write render jobs (Python port of nadir-infer.js decode)."""
3
+ import argparse, gzip, json, math, pickle
4
+ from pathlib import Path
5
+ import numpy as np
6
+ import torch
7
+
8
+ from model import NadirNet
9
+ from train import collate, prep
10
+ from mkjobs import job_from
11
+
12
+ TREAT = ["keep", "dim", "invert", "light", "plate"]
13
+
14
+
15
+ from safety import oklab, lab_to_rgb, contrast, safe_treatment, enforce_contrast, is_media_keep, hue_guard, keep_invisible
16
+
17
+
18
+ def decode(L, lab, alpha_logit, img_logits, rels):
19
+ toks = L["tokens"][:256]
20
+ out, keep = [], set()
21
+ for i, t in enumerate(toks):
22
+ ain = min(1 - 1e-3, max(1e-3, t["rgba"][3]))
23
+ alpha = 1 / (1 + math.exp(-(alpha_logit[i] + math.log(ain / (1 - ain)))))
24
+ a = 1.0 if (t["rgba"][3] >= 0.999 and alpha > 0.97) else alpha
25
+ out.append({"rgb": lab_to_rgb(*lab[i]), "a": a, "L": float(lab[i][0])})
26
+ if is_media_keep(t): # shown as-is: later safety passes see its real colour
27
+ keep.add(i)
28
+ rgb = [int(x) for x in t["rgba"][:3]]
29
+ out[i] = {"rgb": rgb, "a": float(t["rgba"][3]), "L": oklab(*rgb)[0]}
30
+ else:
31
+ hue_guard(t, out[i])
32
+ for i, t in enumerate(toks):
33
+ if t["f"][24] > 0.5 and t["role"] == 0 and out[i]["L"] > 0.35:
34
+ 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
35
+ keep_invisible(toks, out, rels, keep)
36
+ enforce_contrast(toks, out, rels, keep)
37
+ def hx(o):
38
+ h = "#%02x%02x%02x" % tuple(o["rgb"])
39
+ return h + ("%02x" % round(o["a"] * 255) if o["a"] < 0.999 else "")
40
+ colors = [None if i in keep else hx(o) for i, o in enumerate(out)]
41
+ treat = {}
42
+ for j, im in enumerate(L["images"][:32]):
43
+ t = TREAT[int(np.argmax(img_logits[j]))] if j < len(img_logits) else "keep"
44
+ u = im.get("underTok", -1)
45
+ treat[str(j)] = safe_treatment(t, im, out[u]["L"] if u is not None and 0 <= u < len(out) else None)
46
+ return colors, sorted(keep), treat
47
+
48
+
49
+ def main():
50
+ ap = argparse.ArgumentParser()
51
+ ap.add_argument("--ckpt", required=True)
52
+ ap.add_argument("--data", required=True)
53
+ ap.add_argument("--crawl", required=True)
54
+ ap.add_argument("--out", required=True)
55
+ ap.add_argument("--n", type=int, default=150)
56
+ ap.add_argument("--src", type=int, default=0)
57
+ ap.add_argument("--tag", default="nadir")
58
+ ap.add_argument("--also", default="haiku,dr", help="extra render tags to emit for the same pages")
59
+ ap.add_argument("--haiku", default="")
60
+ args = ap.parse_args()
61
+ ck = torch.load(args.ckpt, map_location="cpu")
62
+ model = NadirNet(**ck["cfg"]); model.load_state_dict(ck["model"]); model.eval()
63
+ dev = "cuda" if torch.cuda.is_available() else "cpu"
64
+ model.to(dev)
65
+ data = pickle.load(open(args.data, "rb"))
66
+ val = [s for s in data["val"] if s["src"] == args.src][: args.n]
67
+ with open(args.out, "w") as f:
68
+ for s in val:
69
+ s = prep(s)
70
+ L = json.load(gzip.open(Path(args.crawl) / "pages" / s["id"] / "light.json.gz"))
71
+ b = collate([s], dev)
72
+ with torch.no_grad():
73
+ 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"])
74
+ n = len(s["role"])
75
+ colors, keep, treat = decode(L, lab[0, :n].float().cpu().numpy().tolist(), al[0, :n].float().cpu().numpy().tolist(),
76
+ il[0].float().cpu().numpy().tolist(), s["rel"])
77
+ colors = colors + [None] * (len(L["tokens"]) - len(colors))
78
+ j = job_from(L, s["id"], args.tag, colors, treat)
79
+ j["keep"] = sorted(set(j["keep"]) | set(keep))
80
+ f.write(json.dumps(j) + "\n")
81
+ if "dr" in args.also:
82
+ f.write(json.dumps(job_from(L, s["id"], "dr", [None] * len(L["tokens"]), {}, baseline=True)) + "\n")
83
+ if "haiku" in args.also and args.haiku:
84
+ hp = Path(args.haiku) / f"{s['id']}.json"
85
+ if hp.exists():
86
+ h = json.loads(hp.read_text())
87
+ outs = [h["colors"].get(str(i)) for i in range(len(L["tokens"]))]
88
+ f.write(json.dumps(job_from(L, s["id"], "haiku", outs, h["images"])) + "\n")
89
+ print("wrote", len(val), "pages")
90
+
91
+
92
+ if __name__ == "__main__":
93
+ main()
training/quantize.py ADDED
@@ -0,0 +1,330 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """Calibrated weight-only quantization of NadirNet into ONNX MatMulNBits (8-bit and 4-bit).
3
+
4
+ Layer-sequential GPTQ: each layer's Hessian comes from real crawled pages flowing through the already
5
+ quantized earlier layers. Each group of 32 weights gets asymmetric uint codes whose clipping range is
6
+ chosen by an MSE search. Codes are packed in ORT's MatMulNBits layout, and every variant is checked
7
+ three ways: against a PyTorch fake-quant copy (packing parity), against fp32 (fidelity) and against
8
+ the Haiku targets (quality)."""
9
+ import argparse, copy, json, os, pickle, time
10
+ import numpy as np
11
+ import torch
12
+ import torch.nn as nn
13
+
14
+ from model import NadirNet
15
+ from train import prep
16
+ from export import feeds_for, fidelity, bench, NAMES_IN
17
+
18
+ GS = 32
19
+
20
+
21
+ class Stop(Exception):
22
+ pass
23
+
24
+
25
+ def group_params(w, bits, maxshrink, grid=100):
26
+ """Per-row scale/zero-point for one group (rows, GS), minimizing |q(w)-w|^2.4 over clipping ranges."""
27
+ maxq = 2 ** bits - 1
28
+ xmin = w.min(1).values.clamp(max=0)
29
+ xmax = w.max(1).values.clamp(min=0)
30
+ best = torch.full_like(xmin, float("inf"))
31
+ bs, bz = torch.ones_like(xmin), torch.zeros_like(xmin)
32
+ for i in range(int(maxshrink * grid) + 1):
33
+ p = 1 - i / grid
34
+ lo, hi = p * xmin, p * xmax
35
+ s = ((hi - lo) / maxq).clamp(min=1e-10)
36
+ z = torch.round(-lo / s).clamp(0, maxq)
37
+ q = (torch.clamp(torch.round(w / s[:, None]) + z[:, None], 0, maxq) - z[:, None]) * s[:, None]
38
+ err = (q - w).abs().pow(2.4).sum(1)
39
+ m = err < best
40
+ best[m], bs[m], bz[m] = err[m], s[m], z[m]
41
+ return bs, bz
42
+
43
+
44
+ def gptq(W, H, bits, percdamp=0.01, blocksize=128, gs=GS):
45
+ """W (out, Kpad) with Kpad % gs == 0, H (Kpad, Kpad); gs=None quantizes per output channel (one group per row).
46
+ Returns dequantized W, codes, scales, zero points."""
47
+ maxq = 2 ** bits - 1
48
+ maxshrink = 0.5 if bits == 4 else 0.2
49
+ W = W.clone().double()
50
+ H = H.clone().double()
51
+ rows, cols = W.shape
52
+ dead = torch.diag(H) == 0
53
+ H[dead, dead] = 1
54
+ W[:, dead] = 0
55
+ H += percdamp * torch.mean(torch.diag(H)) * torch.eye(cols, dtype=H.dtype)
56
+ Hinv = torch.linalg.cholesky(torch.cholesky_inverse(torch.linalg.cholesky(H)), upper=True)
57
+ Q = torch.zeros_like(W)
58
+ codes = torch.zeros(rows, cols, dtype=torch.uint8)
59
+ if gs is None: # per channel: one range per output row, chosen on the whole row up front
60
+ s, z = group_params(W, bits, maxshrink)
61
+ S, Z = s[:, None].clone(), z[:, None].clone()
62
+ else:
63
+ S = torch.zeros(rows, cols // gs, dtype=torch.float64)
64
+ Z = torch.zeros(rows, cols // gs, dtype=torch.float64)
65
+ for i1 in range(0, cols, blocksize):
66
+ i2 = min(i1 + blocksize, cols)
67
+ W1 = W[:, i1:i2].clone()
68
+ E1 = torch.zeros_like(W1)
69
+ H1 = Hinv[i1:i2, i1:i2]
70
+ for i in range(i2 - i1):
71
+ c = i1 + i
72
+ if gs is not None and c % gs == 0: # groups never straddle blocks (blocksize % gs == 0)
73
+ s, z = group_params(W1[:, i:i + gs], bits, maxshrink)
74
+ S[:, c // gs], Z[:, c // gs] = s, z
75
+ w = W1[:, i]
76
+ qc = torch.clamp(torch.round(w / s) + z, 0, maxq)
77
+ q = (qc - z) * s
78
+ codes[:, c] = qc.to(torch.uint8)
79
+ Q[:, c] = q
80
+ err = (w - q) / H1[i, i]
81
+ W1[:, i:] -= err[:, None] * H1[i, i:][None, :]
82
+ E1[:, i] = err
83
+ W[:, i2:] -= E1 @ Hinv[i1:i2, i2:]
84
+ return Q.float(), codes, S.float(), Z.to(torch.uint8)
85
+
86
+
87
+ def pack(codes, S, Z, bits):
88
+ """MatMulNBits layout: B (N, k_blocks, blob), scales (N, k_blocks), zero points (N, ceil(k_blocks*bits/8))."""
89
+ c = codes.numpy()
90
+ N, Kp = c.shape
91
+ kb = Kp // GS
92
+ z = Z.numpy().astype(np.uint8)
93
+ if bits == 8:
94
+ B = c.reshape(N, kb, GS).astype(np.uint8)
95
+ zp = z
96
+ else:
97
+ c = c.reshape(N, kb, GS)
98
+ B = (c[..., 0::2] | (c[..., 1::2] << 4)).astype(np.uint8)
99
+ if kb % 2:
100
+ z = np.concatenate([z, np.full((N, 1), 8, np.uint8)], 1)
101
+ zp = (z[:, 0::2] | (z[:, 1::2] << 4)).astype(np.uint8)
102
+ return B, S.numpy().astype(np.float32), zp
103
+
104
+
105
+ def linear_order(model, sample):
106
+ """Linear modules in the order the forward pass calls them."""
107
+ order, hooks = [], []
108
+ for name, mod in model.named_modules():
109
+ if isinstance(mod, nn.Linear):
110
+ hooks.append(mod.register_forward_pre_hook(lambda m, a, n=name: order.append(n) if n not in order else None))
111
+ with torch.no_grad():
112
+ model(*[torch.from_numpy(sample[k]) for k in NAMES_IN])
113
+ for h in hooks:
114
+ h.remove()
115
+ return order
116
+
117
+
118
+ def onnx_matmuls(model, onnx_model):
119
+ """Map each Linear to the MatMul node whose initializer is its transposed weight (Gemm exports stay fp32)."""
120
+ from onnx import numpy_helper
121
+ inits = {i.name: i for i in onnx_model.graph.initializer}
122
+ by_shape = {}
123
+ for node in onnx_model.graph.node:
124
+ if node.op_type == "MatMul" and node.input[1] in inits:
125
+ arr = numpy_helper.to_array(inits[node.input[1]])
126
+ by_shape.setdefault(arr.shape, []).append((node, arr))
127
+ out = {}
128
+ for name, mod in model.named_modules():
129
+ if isinstance(mod, nn.Linear):
130
+ wt = mod.weight.detach().numpy().T
131
+ for node, arr in by_shape.get(wt.shape, []):
132
+ if np.array_equal(arr, wt):
133
+ out[name] = node
134
+ break
135
+ return out
136
+
137
+
138
+ def hessian(model, name, calib):
139
+ """2/n * X^T X of the inputs to one Linear, over all real token rows of the calibration pages."""
140
+ mod = dict(model.named_modules())[name]
141
+ K = mod.in_features
142
+ H = torch.zeros(K, K, dtype=torch.float64)
143
+ n = [0]
144
+
145
+ def hook(m, a):
146
+ x = a[0].detach().reshape(-1, K).double()
147
+ H.add_(x.T @ x)
148
+ n[0] += x.shape[0]
149
+ raise Stop
150
+
151
+ h = mod.register_forward_pre_hook(hook)
152
+ with torch.no_grad():
153
+ for f in calib:
154
+ try:
155
+ model(*[torch.from_numpy(f[k]) for k in NAMES_IN])
156
+ except Stop:
157
+ pass
158
+ h.remove()
159
+ return H * (2.0 / max(1, n[0]))
160
+
161
+
162
+ def quantize_model(model, calib, bits_for, order):
163
+ """Sequential GPTQ over `order`; returns the fake-quant model and packed tensors per layer."""
164
+ qm = copy.deepcopy(model).eval()
165
+ mods = dict(qm.named_modules())
166
+ packed = {}
167
+ for name in order:
168
+ bits = bits_for.get(name)
169
+ if not bits:
170
+ continue
171
+ mod = mods[name]
172
+ K = mod.in_features
173
+ H = hessian(qm, name, calib)
174
+ t = time.time()
175
+ # blocks of 32 along K (zero-padded to a multiple of 32)
176
+ Kp = -(-K // GS) * GS
177
+ Hp = torch.zeros(Kp, Kp, dtype=torch.float64); Hp[:K, :K] = H
178
+ Wp = torch.zeros(mod.out_features, Kp); Wp[:, :K] = mod.weight.detach()
179
+ Q, codes, S, Z = gptq(Wp, Hp, bits)
180
+ if bits == 8:
181
+ # 8-bit: uint8 codes as a (K/32, 32, N) tensor with per-block scale and zero point, dequantized in-graph
182
+ N, kb = mod.out_features, Kp // GS
183
+ packed[name] = ("dq8", K, codes.numpy().T.reshape(kb, GS, N).copy(), S.numpy().T[:, None, :].astype(np.float32).copy(),
184
+ Z.numpy().T[:, None, :].astype(np.uint8).copy())
185
+ else:
186
+ # 4-bit: packed as MatMulNBits
187
+ packed[name] = ("nbits", K) + pack(codes, S, Z, bits)
188
+ with torch.no_grad():
189
+ mod.weight.copy_(Q[:, :K])
190
+ print(f" {name} {tuple(mod.weight.shape)} {bits}-bit {time.time() - t:.1f}s", flush=True)
191
+ return qm, packed
192
+
193
+
194
+ def build_onnx(fp32_path, mapping, packed, out_path, accuracy_level=0):
195
+ import onnx
196
+ from onnx import helper, numpy_helper
197
+ m = onnx.load(fp32_path)
198
+ g = m.graph
199
+ drop_inits = set()
200
+ for name, (kind, K, B, S, zp) in packed.items():
201
+ node = mapping[name]
202
+ wname = node.input[1]
203
+ idx = list(g.node).index(node)
204
+ if kind == "dq8":
205
+ # uint8 codes dequantized with plain Cast/Sub/Mul: every onnxruntime constant-folds this into exact fp32
206
+ # weights at load. (DequantizeLinear->MatMul gets rewritten by onnxruntime into an activation-quantized
207
+ # kernel, which costs accuracy.)
208
+ kb, gs, N = B.shape
209
+ inits = [numpy_helper.from_array(B, wname + "_q8"), numpy_helper.from_array(zp, wname + "_zero_point"),
210
+ numpy_helper.from_array(S, wname + "_scale"), numpy_helper.from_array(np.array([kb * gs, N], np.int64), wname + "_kn")]
211
+ g.initializer.extend(inits)
212
+ base = (node.name or wname) + "_deq8"
213
+ F = onnx.TensorProto.FLOAT
214
+ deq = [helper.make_node("Cast", [inits[0].name], [wname + "_f"], name=base + "_cast", to=F),
215
+ helper.make_node("Cast", [inits[1].name], [wname + "_zf"], name=base + "_castz", to=F),
216
+ helper.make_node("Sub", [wname + "_f", wname + "_zf"], [wname + "_c"], name=base + "_sub"),
217
+ helper.make_node("Mul", [wname + "_c", inits[2].name], [wname + "_d3"], name=base + "_mul"),
218
+ helper.make_node("Reshape", [wname + "_d3", inits[3].name], [wname + ("_dq" if kb * gs == K else "_dp")], name=base + "_reshape")]
219
+ if kb * gs != K: # drop the zero padding rows
220
+ lims = [numpy_helper.from_array(np.array(v, np.int64), wname + n) for v, n in (([0], "_s0"), ([K], "_sK"), ([0], "_ax"))]
221
+ g.initializer.extend(lims)
222
+ deq.append(helper.make_node("Slice", [wname + "_dp"] + [i.name for i in lims], [wname + "_dq"], name=base + "_slice"))
223
+ g.node[idx].input[1] = wname + "_dq" # `node` belongs to another copy of the graph: edit this one
224
+ for k, dn in enumerate(deq):
225
+ g.node.insert(idx + k, dn)
226
+ else:
227
+ N = B.shape[0]
228
+ inits = [numpy_helper.from_array(B, wname + "_Q4"), numpy_helper.from_array(S, wname + "_scales"),
229
+ numpy_helper.from_array(zp, wname + "_zero_points")]
230
+ g.initializer.extend(inits)
231
+ kw = dict(K=K, N=N, bits=4, block_size=GS)
232
+ if accuracy_level:
233
+ kw["accuracy_level"] = accuracy_level
234
+ new = helper.make_node("MatMulNBits", [node.input[0]] + [i.name for i in inits], list(node.output),
235
+ name=(node.name or wname) + "_Q4", domain="com.microsoft", **kw)
236
+ g.node.remove(node)
237
+ g.node.insert(idx, new)
238
+ drop_inits.add(wname)
239
+ keep = [i for i in g.initializer if i.name not in drop_inits]
240
+ del g.initializer[:]
241
+ g.initializer.extend(keep)
242
+ if any(n.domain == "com.microsoft" for n in g.node) and not any(o.domain == "com.microsoft" for o in m.opset_import):
243
+ m.opset_import.append(helper.make_opsetid("com.microsoft", 1))
244
+ onnx.checker.check_model(m)
245
+ onnx.save(m, out_path)
246
+
247
+
248
+ def quality(sess, samples):
249
+ """dE vs the training targets (Haiku for src 0, native dark for src 1) and image-treatment accuracy."""
250
+ agg = {0: [0.0, 0.0], 1: [0.0, 0.0]}
251
+ ok = n = 0
252
+ for s in samples:
253
+ f = feeds_for(s)
254
+ lab, _, imgl = sess.run(None, f)
255
+ k = len(s["role"])
256
+ m = s["tmask"][:k]
257
+ if m.any():
258
+ d = np.linalg.norm(lab[0, :k] - s["lab_t"][:k], axis=-1)
259
+ w = s["w"][:k] * m
260
+ agg[int(s["src"])][0] += float((d * w).sum() / max(w.sum(), 1e-9)); agg[int(s["src"])][1] += 1
261
+ mi = len(s["IF"])
262
+ if mi:
263
+ t = np.asarray(s["imgt"][:mi])
264
+ v = t >= 0
265
+ ok += int((imgl[0, :mi].argmax(-1)[v] == t[v]).sum()); n += int(v.sum())
266
+ 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)}
267
+
268
+
269
+ def main():
270
+ ap = argparse.ArgumentParser()
271
+ ap.add_argument("--ckpt", required=True)
272
+ ap.add_argument("--data", required=True)
273
+ ap.add_argument("--fp32", required=True, help="fp32 ONNX from export.py")
274
+ ap.add_argument("--out", required=True)
275
+ ap.add_argument("--calib", type=int, default=256)
276
+ ap.add_argument("--eval", type=int, default=400)
277
+ ap.add_argument("--threads", type=int, default=8)
278
+ args = ap.parse_args()
279
+ torch.set_num_threads(args.threads)
280
+ os.makedirs(args.out, exist_ok=True)
281
+ import onnx, onnxruntime as ort
282
+ ck = torch.load(args.ckpt, map_location="cpu")
283
+ model = NadirNet(**ck["cfg"]); model.load_state_dict(ck["model"]); model.eval()
284
+ data = pickle.load(open(args.data, "rb"))
285
+ hv = [s for s in data["val"] if s["src"] == 0]; nv = [s for s in data["val"] if s["src"] == 1]
286
+ val = [prep(s) for s in hv[: args.eval * 3 // 4] + nv[: args.eval - min(len(hv), args.eval * 3 // 4)]] # both sources
287
+ rng = np.random.default_rng(0)
288
+ pick = rng.choice(len(data["train"]), size=min(args.calib, len(data["train"])), replace=False)
289
+ calib = [feeds_for(prep(data["train"][i])) for i in pick]
290
+ order = linear_order(model, calib[0])
291
+ mapping = onnx_matmuls(model, onnx.load(args.fp32))
292
+ quantizable = [n for n in order if n in mapping and not n.endswith(".rel")]
293
+ print(f"{len(quantizable)}/{len(order)} linears map to MatMul; fp32 kept: {[n for n in order if n not in quantizable]}", flush=True)
294
+ edges = [n for n in quantizable if n.startswith(("tok_in.", "img_in.", "color_head.", "img_head."))]
295
+ plans = {"int8": {n: 8 for n in quantizable},
296
+ "int4": {n: (8 if n in edges else 4) for n in quantizable}}
297
+ so = ort.SessionOptions(); so.intra_op_num_threads = 4
298
+ ref = ort.InferenceSession(args.fp32, so, providers=["CPUExecutionProvider"])
299
+ 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"},
300
+ "group_size": GS, "calib_pages": len(calib), "eval_pages": len(val), "fp32_layers": [n for n in order if n not in quantizable],
301
+ "int4_8bit_edge_layers": edges, "variants": {}}
302
+ r = quality(ref, val); r["size_mb"] = round(os.path.getsize(args.fp32) / 1e6, 2); r["cpu_ms_median"] = round(bench(ref, val), 2)
303
+ report["variants"]["fp32"] = r
304
+ print("fp32", r, flush=True)
305
+ for vname, bits_for in plans.items():
306
+ print(f"== {vname}", flush=True)
307
+ qm, packed = quantize_model(model, calib, bits_for, quantizable)
308
+ path = os.path.join(args.out, f"nadir-{vname}.onnx")
309
+ build_onnx(args.fp32, mapping, packed, path)
310
+ sess = ort.InferenceSession(path, so, providers=["CPUExecutionProvider"])
311
+ # packing parity: ORT MatMulNBits graph vs the PyTorch fake-quant model
312
+ par = 0.0
313
+ with torch.no_grad():
314
+ for s in val[:8]:
315
+ f = feeds_for(s)
316
+ tq = qm(*[torch.from_numpy(f[k]) for k in NAMES_IN])[0].numpy()
317
+ par = max(par, float(np.abs(tq - sess.run(None, f)[0]).max()))
318
+ r = fidelity(sess, ref, val)
319
+ r.update(quality(sess, val))
320
+ r["pack_parity_max_abs"] = par
321
+ r["size_mb"] = round(os.path.getsize(path) / 1e6, 2)
322
+ r["cpu_ms_median"] = round(bench(sess, val), 2)
323
+ report["variants"][vname] = r
324
+ print(vname, r, flush=True)
325
+ torch.save({"cfg": ck["cfg"], "model": qm.state_dict()}, os.path.join(args.out, f"fakequant-{vname}.pt"))
326
+ json.dump(report, open(os.path.join(args.out, "quant_report.json"), "w"), indent=1)
327
+
328
+
329
+ if __name__ == "__main__":
330
+ main()
training/safety.py ADDED
@@ -0,0 +1,188 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Deterministic safety rules shared by the dataset builder (target repair) and the decoder (runtime).
2
+ Mirrors nadir-infer.js / nadir-map.js."""
3
+ import math
4
+
5
+
6
+ def lin(u):
7
+ u = u / 255.0
8
+ return u / 12.92 if u <= 0.04045 else ((u + 0.055) / 1.055) ** 2.4
9
+
10
+
11
+ def oklab(r, g, b):
12
+ r, g, b = lin(r), lin(g), lin(b)
13
+ l = (0.4122214708 * r + 0.5363325363 * g + 0.0514459929 * b) ** (1 / 3)
14
+ m = (0.2119034982 * r + 0.6806995451 * g + 0.1073969566 * b) ** (1 / 3)
15
+ s = (0.0883024619 * r + 0.2817188376 * g + 0.6299787005 * b) ** (1 / 3)
16
+ return [0.2104542553 * l + 0.7936177850 * m - 0.0040720468 * s, 1.9779984951 * l - 2.4285922050 * m + 0.4505937099 * s,
17
+ 0.0259040371 * l + 0.7827717662 * m - 0.8086757660 * s]
18
+
19
+
20
+ def lab_to_rgb(L, a, b):
21
+ l_, m_, s_ = L + 0.3963377774 * a + 0.2158037573 * b, L - 0.1055613458 * a - 0.0638541728 * b, L - 0.0894841775 * a - 1.2914855480 * b
22
+ l, m, s = l_ ** 3, m_ ** 3, s_ ** 3
23
+ def enc(x):
24
+ x = max(0.0, min(1.0, x))
25
+ return round(255 * (12.92 * x if x <= 0.0031308 else 1.055 * x ** (1 / 2.4) - 0.055))
26
+ return [enc(4.0767416621 * l - 3.3077115913 * m + 0.2309699292 * s), enc(-1.2684380046 * l + 2.6097574011 * m - 0.3413193965 * s),
27
+ enc(-0.0041960863 * l - 0.7034186147 * m + 1.7076147010 * s)]
28
+
29
+
30
+ def rel_lum(c):
31
+ return 0.2126 * lin(c[0]) + 0.7152 * lin(c[1]) + 0.0722 * lin(c[2])
32
+
33
+
34
+ def contrast(a, b):
35
+ x, y = rel_lum(a), rel_lum(b)
36
+ return (max(x, y) + 0.05) / (min(x, y) + 0.05)
37
+
38
+
39
+ def adaptable(im, s, surf_L=None):
40
+ """A logo on its own white rectangle: adapt it instead of leaving a white box. Only images the page marks as its
41
+ logo (logo/brand/wordmark naming, or the link home); product shots, photos and other images are never touched.
42
+ Without pixel stats a logo still tries (the adaptation itself checks for the white background)."""
43
+ if not (im.get("logo") or (im.get("homeLink") and im.get("header"))) or im["h"] > 200 or (surf_L is not None and surf_L >= 0.4):
44
+ return False
45
+ if not s:
46
+ return im["w"] * im["h"] <= 250000
47
+ 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
48
+
49
+
50
+ def is_light_backdrop(im, s):
51
+ """Opaque, light, flat (almost no lightness variation, so never a photo; SVG backgrounds can't be photos, so a
52
+ smooth gradient qualifies too) and big enough to sit behind content. Logos drawn as backgrounds are never cleared."""
53
+ if im.get("kind") != "bg" or im.get("logo") or im.get("homeLink") or not s or s["transp"] >= 0.05 or s["meanL"] <= 0.8 or im["w"] * im["h"] < 40000:
54
+ return False
55
+ svg = im.get("src", "").split("?")[0].lower().endswith(".svg")
56
+ return s["stdL"] < (0.15 if svg else 0.07)
57
+
58
+
59
+ def safe_treatment(t, im, surf_L=None):
60
+ """'adapt:<fallback>' for logos on a white rectangle (fallback used when there is nothing to adapt)."""
61
+ base = _base_treatment(t, im, surf_L)
62
+ return "adapt:" + base if base != "clear" and adaptable(im, im.get("stats"), surf_L) else base
63
+
64
+
65
+ def _base_treatment(t, im, surf_L=None):
66
+ s = im.get("stats")
67
+ # a light, flat CSS background image (texture, pale wash, pattern) sits behind content whose text is now light:
68
+ # clear it so the dark surface designed for that element shows through. Photos and detailed images stay.
69
+ if is_light_backdrop(im, s):
70
+ return "clear"
71
+ # a mostly dark graphic on transparency would vanish on a dark surface: hue-preserving inversion
72
+ 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):
73
+ return "invert"
74
+ if not t or t == "keep":
75
+ return "keep"
76
+ photo = (s["transp"] < 0.02 and s["colors"] > 0.18 and s["stdL"] > 0.12) if s else (im["w"] * im["h"] > 60000)
77
+ if photo and t != "dim":
78
+ return "keep"
79
+ if t == "dim" and im["w"] * im["h"] < 20000:
80
+ return "keep"
81
+ # a transparent graphic that is already light reads fine on a dark surface: leave it alone
82
+ if t in ("plate", "light", "invert") and s and s["transp"] > 0.15 and s["meanL"] >= 0.55 and s.get("black", 0) < 0.15 and (surf_L is None or surf_L < 0.4):
83
+ return "keep"
84
+ # no light boxes on a dark page: a dark logo on transparency gets a hue-preserving inversion instead of a plate,
85
+ # an opaque one already carries its own background
86
+ if t == "plate":
87
+ return "invert" if (s is None or s["transp"] > 0.15) else "keep"
88
+ return t
89
+
90
+
91
+ FG_ROLES = (1, 5, 6) # text, svg fill, svg stroke
92
+
93
+
94
+ def hue_guard(t, o):
95
+ """Brand colours keep their hue: a coloured input whose output drifts more than 45 degrees is turned back to
96
+ its own hue (keeping the predicted lightness and chroma); neutral inputs get at most a faint tint."""
97
+ rgb = [int(x) for x in t["rgba"][:3]]
98
+ i = oklab(*rgb)
99
+ q = oklab(*o["rgb"])
100
+ cin, cout = math.hypot(i[1], i[2]), math.hypot(q[1], q[2])
101
+ if cin > 0.03 and cout > 0.03:
102
+ d = abs((math.atan2(q[2], q[1]) - math.atan2(i[2], i[1]) + math.pi) % (2 * math.pi) - math.pi)
103
+ if d <= math.radians(45):
104
+ return
105
+ h = math.atan2(i[2], i[1])
106
+ o["rgb"] = lab_to_rgb(q[0], cout * math.cos(h), cout * math.sin(h))
107
+ elif cin <= 0.03 and cout > 0.06:
108
+ o["rgb"] = lab_to_rgb(q[0], q[1] * 0.06 / cout, q[2] * 0.06 / cout)
109
+
110
+
111
+ def keep_invisible(toks, out, rels, keep):
112
+ """Things that were invisible against their surface in the light design stay invisible: a gradient stop
113
+ with exactly a surface's colour (text fades, edge fades) takes that surface's dark colour, and a border that
114
+ matched the surface it sits on takes the surface's dark colour too (otherwise: dark slabs, white edge bars)."""
115
+ n = len(toks)
116
+ key = lambda t: tuple(int(x) for x in t["rgba"][:3])
117
+ surf = {}
118
+ for i, t in enumerate(toks): # first (most prominent) opaque surface of each colour
119
+ if t["role"] == 0 and t["rgba"][3] >= 0.999 and i not in keep and out[i]["a"] >= 0.999:
120
+ surf.setdefault(key(t), i)
121
+ for i, t in enumerate(toks):
122
+ if t["role"] == 4 and i not in keep and key(t) in surf:
123
+ j = surf[key(t)]
124
+ out[i]["rgb"], out[i]["L"] = list(out[j]["rgb"]), out[j]["L"]
125
+ best = {}
126
+ for typ, a, b, w in rels:
127
+ a, b = int(a), int(b)
128
+ if int(typ) == 2 and a < n and b < n and toks[a]["role"] == 2 and toks[b]["role"] == 0 and (a not in best or w > best[a][0]):
129
+ best[a] = (w, b)
130
+ for i, (_, j) in best.items():
131
+ if i in keep or j in keep:
132
+ continue
133
+ if contrast(toks[i]["rgba"][:3], toks[j]["rgba"][:3]) < 1.05 and toks[i]["rgba"][3] >= 0.5:
134
+ out[i]["rgb"], out[i]["L"] = list(out[j]["rgb"]), out[j]["L"]
135
+
136
+
137
+ def is_media_keep(t):
138
+ """Colours drawn over video/photos that already read as dark mode stay exactly as they are: dark scrims and
139
+ player bars, light captions and icons. Everything else over media (a white page or a frosted white button
140
+ laid over a decorative canvas) is redesigned like any other colour."""
141
+ f = t["f"]
142
+ if f[48] < 0.5:
143
+ return False
144
+ if t["role"] in (0, 3, 4): # bg, shadow, gradient: dark scrims, or saturated accents such as a progress bar
145
+ return f[24] <= 0.5 and f[11] < 0.25 and (f[0] < 0.5 or f[3] > 0.4)
146
+ return f[0] > 0.6
147
+
148
+
149
+ def enforce_contrast(toks, out, rels, keep):
150
+ """Keep every foreground colour readable on all the surfaces it sits on, never asking for more contrast
151
+ than the light design had. One lightness per token (tokens are shared), chosen on a grid to minimise the
152
+ weighted log-contrast deficit, with a small pull towards the model's own lightness."""
153
+ n = len(toks)
154
+ surf = {}
155
+ for typ, a, b, w in rels:
156
+ a, b = int(a), int(b)
157
+ if int(typ) == 0 and a < n and b < n and toks[b]["role"] == 0 and out[b]["a"] >= 0.999:
158
+ surf.setdefault(a, []).append((b, float(w) + 0.1)) # log-area weights: small buttons still count
159
+ for i, lst in surf.items():
160
+ t = toks[i]
161
+ if i in keep or t["role"] not in FG_ROLES or out[i]["a"] < 0.999:
162
+ continue
163
+ if t["role"] == 1:
164
+ fs, bold = t["f"][21] * 32, t["f"][23] > 0.5
165
+ need0 = 3.0 if (fs >= 24 or (fs >= 18.5 and bold)) else 4.5
166
+ else:
167
+ need0 = 3.0
168
+ reqs = []
169
+ for j, w in lst:
170
+ need = min(need0, contrast(t["rgba"][:3], toks[j]["rgba"][:3]))
171
+ if need >= 1.5:
172
+ reqs.append((out[j]["rgb"], need, w))
173
+ if not reqs:
174
+ continue
175
+ def deficit(rgb):
176
+ return sum(w * max(0.0, math.log(need) - math.log(contrast(rgb, bg))) for bg, need, w in reqs)
177
+ if deficit(out[i]["rgb"]) <= 0:
178
+ continue
179
+ o = oklab(*out[i]["rgb"])
180
+ W = sum(w for _, _, w in reqs)
181
+ best = None
182
+ for k in range(51):
183
+ Lx = k / 50
184
+ rgb = lab_to_rgb(Lx, o[1], o[2])
185
+ c = deficit(rgb) / W + 0.05 * abs(Lx - o[0])
186
+ if best is None or c < best[0]:
187
+ best = (c, rgb, Lx)
188
+ out[i]["rgb"], out[i]["L"] = best[1], best[2]
training/train.py ADDED
@@ -0,0 +1,292 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """Train NadirNet on Haiku-designed and native dark-mode targets."""
3
+ import argparse, json, math, os, pickle, random, time
4
+ from collections import defaultdict
5
+ import numpy as np
6
+ import torch
7
+ import torch.nn.functional as F
8
+
9
+ from model import NadirNet, K_REL, contrast_ratio
10
+ from build_ds import oklab
11
+
12
+ 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))
13
+
14
+
15
+ def lab_to_rgb_np(lab):
16
+ M1 = np.array([[1.0, 0.3963377774, 0.2158037573], [1.0, -0.1055613458, -0.0638541728], [1.0, -0.0894841775, -1.2914855480]])
17
+ M2 = np.array([[4.0767416621, -3.3077115913, 0.2309699292], [-1.2684380046, 2.6097574011, -0.3413193965], [-0.0041960863, -0.7034186147, 1.7076147010]])
18
+ lin = np.clip((lab @ M1.T) ** 3 @ M2.T, 0, 1)
19
+ return np.where(lin <= 0.0031308, 12.92 * lin, 1.055 * np.power(lin, 1 / 2.4) - 0.055)
20
+
21
+
22
+ def prep(s):
23
+ """Precompute OKLab of inputs/targets once per sample."""
24
+ if "lab_in" in s:
25
+ return s
26
+ s["lab_in"] = oklab(s["rgba"][:, :3]).astype(np.float32)
27
+ t = s["tgt"]
28
+ ok = np.isfinite(t[:, 0])
29
+ lab_t = np.zeros((len(t), 3), np.float32)
30
+ if ok.any():
31
+ lab_t[ok] = oklab(t[ok, :3]).astype(np.float32)
32
+ s["lab_t"], s["tmask"] = lab_t, ok
33
+ s["alpha_t"] = np.where(ok, t[:, 3], 1).astype(np.float32)
34
+ km = np.zeros(len(t), bool)
35
+ if len(s["keep"]):
36
+ km[s["keep"]] = True
37
+ s["tmask"] = s["tmask"] & ~km
38
+ return s
39
+
40
+
41
+ def augment(s, rng):
42
+ s = dict(s)
43
+ n = len(s["role"])
44
+ # token dropout (simulates ad blockers, A/B variants, partial loads)
45
+ if rng.random() < 0.6 and n > 8:
46
+ p = rng.uniform(0, 0.25)
47
+ root = s["F"][:, 24] > 0.5
48
+ keepm = (rng.random(n) > p) | root
49
+ idx = np.nonzero(keepm)[0]
50
+ remap = -np.ones(n, np.int64); remap[idx] = np.arange(len(idx))
51
+ for k in ("F", "role", "rgba", "w", "lab_in", "lab_t", "tmask", "alpha_t"):
52
+ s[k] = s[k][idx]
53
+ rel = s["rel"]
54
+ if len(rel):
55
+ a, b = rel[:, 1].astype(int), rel[:, 2].astype(int)
56
+ good = (remap[a] >= 0) & (remap[b] >= 0)
57
+ rel = rel[good].copy(); rel[:, 1] = remap[a[good]]; rel[:, 2] = remap[b[good]]
58
+ s["rel"] = rel
59
+ iu = s["img_under"].copy(); m = iu >= 0; iu[m] = remap[iu[m]]; s["img_under"] = iu
60
+ n = len(idx)
61
+ # hue rotation of the whole design (inputs and targets together)
62
+ if rng.random() < 0.3:
63
+ th = rng.uniform(0, 2 * math.pi)
64
+ c, si = math.cos(th), math.sin(th)
65
+ rot = lambda a, b: (c * a - si * b, si * a + c * b)
66
+ F_ = s["F"].copy()
67
+ for (ia, ib) in ((1, 2), (35, 36), (40, 41), (5, 6)):
68
+ F_[:, ia], F_[:, ib] = rot(F_[:, ia], F_[:, ib])
69
+ lab = s["lab_in"].copy(); lab[:, 1], lab[:, 2] = rot(lab[:, 1], lab[:, 2])
70
+ F_[:, 7:10] = lab_to_rgb_np(lab.astype(np.float64)).astype(np.float32)
71
+ s["F"], s["lab_in"] = F_, lab.astype(np.float32)
72
+ lt = s["lab_t"].copy(); lt[:, 1], lt[:, 2] = rot(lt[:, 1], lt[:, 2]); s["lab_t"] = lt.astype(np.float32)
73
+ return s
74
+
75
+
76
+ def collate(batch, device):
77
+ B = len(batch)
78
+ N = max(1, max(len(s["role"]) for s in batch))
79
+ M = max(len(s["IF"]) for s in batch)
80
+ T = 1 + N + M
81
+ tok_f = np.zeros((B, N, 49), np.float32); role = np.zeros((B, N), np.int64); lab = np.zeros((B, N, 3), np.float32)
82
+ lab_t = np.zeros((B, N, 3), np.float32); tmask = np.zeros((B, N), bool); w = np.zeros((B, N), np.float32)
83
+ alpha_in = np.ones((B, N), np.float32); alpha_t = np.ones((B, N), np.float32)
84
+ 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)
85
+ page = np.zeros((B, 8), np.float32); src = np.zeros(B, np.int64)
86
+ rel = np.zeros((B, T, T, K_REL), np.float32); mask = np.zeros((B, T), bool)
87
+ tb_pair = -np.ones((B, N), np.int64)
88
+ for i, s in enumerate(batch):
89
+ n, m = len(s["role"]), len(s["IF"])
90
+ tok_f[i, :n] = s["F"]; role[i, :n] = s["role"]; lab[i, :n] = s["lab_in"]; lab_t[i, :n] = s["lab_t"]; tmask[i, :n] = s["tmask"]
91
+ w[i, :n] = s["w"]; alpha_in[i, :n] = s["rgba"][:, 3]; alpha_t[i, :n] = s["alpha_t"]
92
+ if m:
93
+ img_f[i, :m] = s["IF"]; thumbs[i, :m] = s["thumbs"].transpose(0, 3, 1, 2) / 255.0; imgt[i, :m] = s["imgt"]
94
+ page[i] = s["page"]; src[i] = s["src"]
95
+ mask[i, :1 + n] = True; mask[i, 1 + N:1 + N + m] = True
96
+ # relations
97
+ r = s["rel"]
98
+ if len(r):
99
+ typ, a, b, wt = r[:, 0].astype(int), r[:, 1].astype(int) + 1, r[:, 2].astype(int) + 1, r[:, 3]
100
+ R = np.zeros((4, n + 1, n + 1), np.float32)
101
+ np.maximum.at(R, (typ, a, b), wt)
102
+ R /= R.max(axis=2, keepdims=True) + 1e-6
103
+ rel[i, :n + 1, :n + 1, 0:4] = R.transpose(1, 2, 0)
104
+ rel[i, :n + 1, :n + 1, 4:8] = R.transpose(2, 1, 0)
105
+ # dominant surface for each text token (contrast loss)
106
+ tbm = typ == 0
107
+ best = {}
108
+ for aa, bb, ww in zip(a[tbm] - 1, b[tbm] - 1, wt[tbm]):
109
+ if ww > best.get(aa, (-1, -1))[0]:
110
+ best[aa] = (ww, bb)
111
+ for aa, (ww, bb) in best.items():
112
+ tb_pair[i, aa] = bb
113
+ for j in range(m):
114
+ u = s["img_under"][j]
115
+ if u >= 0:
116
+ rel[i, 1 + N + j, 1 + u, 8] = 1.0
117
+ rel[i, 1 + u, 1 + N + j, 9] = 1.0
118
+ # same colour, different role
119
+ key = s["rgba"][:, :3] @ np.array([65536, 256, 1], np.float32)
120
+ same = (key[:, None] == key[None, :]) & ~np.eye(n, dtype=bool)
121
+ rel[i, 1:1 + n, 1:1 + n, 10] = same
122
+ if device is None:
123
+ 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,
124
+ img_f=img_f, thumbs=thumbs, imgt=imgt, page=page, src=src, rel=rel, mask=mask, tb_pair=tb_pair)
125
+ t = lambda x: torch.from_numpy(x).to(device, non_blocking=True)
126
+ 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),
127
+ 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))
128
+
129
+
130
+ class Batches(torch.utils.data.Dataset):
131
+ """Each item is a whole (augmented) batch, assembled in a worker process."""
132
+ def __init__(self, data, order, bs, aug, seed):
133
+ self.data, self.order, self.bs, self.aug, self.seed = data, order, bs, aug, seed
134
+
135
+ def __len__(self):
136
+ return math.ceil(len(self.order) / self.bs)
137
+
138
+ def __getitem__(self, k):
139
+ rng = np.random.default_rng((self.seed, k))
140
+ items = [self.data[i] for i in self.order[k * self.bs:(k + 1) * self.bs]]
141
+ if self.aug:
142
+ items = [augment(s, rng) for s in items]
143
+ return collate(items, None)
144
+
145
+
146
+ def to_dev(b, device):
147
+ return {k: (v if torch.is_tensor(v) else torch.from_numpy(v)).to(device, non_blocking=True) for k, v in b.items()}
148
+
149
+
150
+ def huber(x, d):
151
+ ax = x.abs()
152
+ return torch.where(ax < d, 0.5 * x * x / d, ax - 0.5 * d)
153
+
154
+
155
+ def losses(model, b, img_w=None):
156
+ 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"])
157
+ lab = lab.float(); alpha_logit = alpha_logit.float(); img_logits = img_logits.float()
158
+ m = b["tmask"].float() * b["w"]
159
+ d = lab - b["lab_t"]
160
+ per = 2.0 * huber(d[..., 0], 0.02) + huber(d[..., 1], 0.01) + huber(d[..., 2], 0.01)
161
+ lc = (per * m).sum() / m.sum().clamp_min(1)
162
+ # alpha (only where translucency is involved)
163
+ am = m * ((b["alpha_in"] < 0.999) | (b["alpha_t"] < 0.999)).float()
164
+ a_in = b["alpha_in"].clamp(1e-3, 1 - 1e-3)
165
+ alpha = torch.sigmoid(alpha_logit + torch.log(a_in / (1 - a_in)))
166
+ la = (huber(alpha - b["alpha_t"], 0.05) * am).sum() / am.sum().clamp_min(1)
167
+ # readability: text must not lose contrast against its own surface
168
+ pair = b["tb_pair"]
169
+ pm = (pair >= 0) & b["tmask"] & (b["role"] == 1) & (b["alpha_in"] > 0.99)
170
+ if pm.any():
171
+ idx = pair.clamp_min(0)
172
+ bg_pred = torch.gather(lab, 1, idx[..., None].expand(-1, -1, 3))
173
+ bg_t = torch.gather(b["lab_t"], 1, idx[..., None].expand(-1, -1, 3))
174
+ bgm = torch.gather(b["tmask"], 1, idx)
175
+ pm = pm & bgm
176
+ c_pred = contrast_ratio(lab, bg_pred)
177
+ c_t = contrast_ratio(b["lab_t"], bg_t).clamp(max=4.5)
178
+ lcon = (F.relu(c_t - c_pred) * pm.float()).sum() / pm.float().sum().clamp_min(1)
179
+ else:
180
+ lcon = lab.new_zeros(())
181
+ # images
182
+ im = b["imgt"] >= 0
183
+ if im.any():
184
+ li = F.cross_entropy(img_logits[im], b["imgt"][im], weight=img_w, label_smoothing=0.05)
185
+ acc = (img_logits[im].argmax(-1) == b["imgt"][im]).float().mean()
186
+ else:
187
+ li = lab.new_zeros(()); acc = lab.new_zeros(())
188
+ with torch.no_grad():
189
+ de = (d.norm(dim=-1) * m).sum() / m.sum().clamp_min(1)
190
+ dl = (d[..., 0].abs() * m).sum() / m.sum().clamp_min(1)
191
+ return lc, la, lcon, li, dict(dE=de.item(), dL=dl.item(), img_acc=acc.item(), n_img=int(im.sum()))
192
+
193
+
194
+ def evaluate(model, val, device, bs, img_w):
195
+ model.eval()
196
+ agg = defaultdict(float); cnt = defaultdict(float)
197
+ with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
198
+ for src in (0, 1):
199
+ vs = [s for s in val if s["src"] == src]
200
+ for k in range(0, len(vs), bs):
201
+ b = collate(vs[k:k + bs], device)
202
+ lc, la, lcon, li, st = losses(model, b, img_w)
203
+ n = len(vs[k:k + bs])
204
+ agg[f"s{src}_dE"] += st["dE"] * n; agg[f"s{src}_dL"] += st["dL"] * n; agg[f"s{src}_con"] += lcon.item() * n
205
+ cnt[f"s{src}"] += n
206
+ if st["n_img"]:
207
+ agg["img_acc"] += st["img_acc"] * st["n_img"]; cnt["img"] += st["n_img"]
208
+ out = {}
209
+ for k, v in agg.items():
210
+ if k == "img_acc":
211
+ out[k] = v / max(1, cnt["img"])
212
+ else:
213
+ out[k] = v / max(1, cnt[k.split("_")[0]])
214
+ model.train()
215
+ return out
216
+
217
+
218
+ def main():
219
+ ap = argparse.ArgumentParser()
220
+ ap.add_argument("--data", required=True)
221
+ ap.add_argument("--out", required=True)
222
+ ap.add_argument("--d", type=int, default=384)
223
+ ap.add_argument("--layers", type=int, default=8)
224
+ ap.add_argument("--heads", type=int, default=6)
225
+ ap.add_argument("--epochs", type=int, default=40)
226
+ ap.add_argument("--bs", type=int, default=48)
227
+ ap.add_argument("--lr", type=float, default=4e-4)
228
+ ap.add_argument("--wd", type=float, default=0.05)
229
+ ap.add_argument("--drop", type=float, default=0.1)
230
+ ap.add_argument("--native-weight", type=float, default=0.5)
231
+ ap.add_argument("--seed", type=int, default=0)
232
+ ap.add_argument("--workers", type=int, default=8)
233
+ args = ap.parse_args()
234
+ torch.manual_seed(args.seed); rng = np.random.default_rng(args.seed)
235
+ os.makedirs(args.out, exist_ok=True)
236
+ data = pickle.load(open(args.data, "rb"))
237
+ train, val = [prep(s) for s in data["train"]], [prep(s) for s in data["val"]]
238
+ 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)
239
+ device = "cuda"
240
+ model = NadirNet(d=args.d, layers=args.layers, heads=args.heads, drop=args.drop).to(device)
241
+ nparam = sum(p.numel() for p in model.parameters())
242
+ print(f"params {nparam / 1e6:.2f}M", flush=True)
243
+ # image class weights (inverse sqrt frequency)
244
+ cnt = np.ones(5)
245
+ for s in train:
246
+ for v in s["imgt"]:
247
+ if v >= 0:
248
+ cnt[v] += 1
249
+ img_w = torch.tensor((cnt.sum() / cnt) ** 0.5, dtype=torch.float32, device=device)
250
+ img_w = img_w / img_w.mean()
251
+ print("image label counts", cnt.astype(int).tolist(), flush=True)
252
+ opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.wd, betas=(0.9, 0.98))
253
+ steps_per = math.ceil(len(train) / args.bs)
254
+ total = steps_per * args.epochs
255
+ warm = min(800, total // 20)
256
+ sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda st: min(1, (st + 1) / warm) * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * min(1, st / total)))))
257
+ # sample weights: native pages count less (different designers, noisier pairing)
258
+ sw = np.array([1.0 if s["src"] == 0 else args.native_weight for s in train]); sw /= sw.sum()
259
+ best = 1e9; step = 0; log = []
260
+ t0 = time.time()
261
+ for ep in range(args.epochs):
262
+ order = rng.choice(len(train), size=len(train), p=sw)
263
+ run = defaultdict(float); nb = 0
264
+ loader = torch.utils.data.DataLoader(Batches(train, order, args.bs, True, args.seed * 1000 + ep), batch_size=None, shuffle=False,
265
+ num_workers=args.workers, persistent_workers=False, prefetch_factor=4 if args.workers else None)
266
+ for bnp in loader:
267
+ b = to_dev(bnp, device)
268
+ with torch.autocast("cuda", dtype=torch.bfloat16):
269
+ lc, la, lcon, li, st = losses(model, b, img_w)
270
+ loss = lc + 0.5 * la + 0.02 * lcon + 0.05 * li
271
+ opt.zero_grad(set_to_none=True)
272
+ if not torch.isfinite(loss): # one bad page must never poison the weights
273
+ run["skipped"] += 1
274
+ continue
275
+ loss.backward()
276
+ torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
277
+ opt.step(); sched.step(); step += 1
278
+ run["loss"] += loss.item(); run["dE"] += st["dE"]; run["con"] += lcon.item(); run["img"] += li.item(); nb += 1
279
+ ev = evaluate(model, val, device, 64, img_w)
280
+ 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()}}
281
+ log.append(rec)
282
+ print(json.dumps(rec), flush=True)
283
+ 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))
284
+ if score < best:
285
+ best = score
286
+ torch.save({"model": model.state_dict(), "cfg": model.cfg, "epoch": ep, "eval": ev}, os.path.join(args.out, "best.pt"))
287
+ torch.save({"model": model.state_dict(), "cfg": model.cfg, "epoch": args.epochs - 1}, os.path.join(args.out, "last.pt"))
288
+ json.dump({"args": vars(args), "params": nparam, "log": log, "best_score": best}, open(os.path.join(args.out, "train_log.json"), "w"), indent=1)
289
+
290
+
291
+ if __name__ == "__main__":
292
+ main()