| """NadirNet: a relation-biased set transformer over a page's colour palette and images. |
| |
| Tokens: [PAGE] + one token per (colour, role) the page uses + one token per image. |
| Attention logits get a learned per-head bias from the page's structure graph (text-on-surface, |
| surface-on-surface, border-on-surface, same-element, image-on-surface, same-colour), so the model |
| reasons about hierarchy and contrast the way a designer would. Outputs a dark colour per token |
| (OKLab, parameterised relative to the input so hue is kept unless the model chooses otherwise), |
| an alpha, and a compact page context. |
| |
| Element head (v2): decides how every image, CSS background, sprite or pseudo-element icon on the page is shown |
| (keep / invert / adapt / remove) and what it is (photo, logo, icon, formula...). It reuses the trunk's image |
| encoder and reads the page context plus the light surface under the element and the dark surface the palette |
| head designed for it, so it can run for any element at any time (late-loaded icons, sprites) in a tiny second |
| graph without re-running the transformer. |
| """ |
| import math |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| FT, FI, FP = 49, 41, 8 |
| N_ROLES = 7 |
| TREAT = ["keep", "invert", "adapt", "remove"] |
| IMG_ROLES = ["photo", "illustration", "diagram", "logo", "icon", "text", "backdrop", "avatar", "ui"] |
| N_TREAT, N_IROLE, N_CTX = len(TREAT), len(IMG_ROLES), 32 |
| K_REL = 11 |
|
|
|
|
| def fourier_lab(lab, n=6): |
| x = torch.stack([lab[..., 0], lab[..., 1] * 4, lab[..., 2] * 4], -1) |
| freqs = (2.0 ** torch.arange(n, device=lab.device, dtype=lab.dtype)) * math.pi |
| ang = x[..., None] * freqs |
| return torch.cat([torch.sin(ang), torch.cos(ang)], -1).flatten(-2) |
|
|
|
|
| class MLP(nn.Module): |
| def __init__(self, i, h, o, drop=0.0): |
| super().__init__() |
| self.net = nn.Sequential(nn.Linear(i, h), nn.GELU(), nn.Dropout(drop), nn.Linear(h, o)) |
|
|
| def forward(self, x): |
| return self.net(x) |
|
|
|
|
| class ThumbCNN(nn.Module): |
| """32x32 RGBA thumbnail -> embedding. Alpha matters: transparent logos behave differently from photos.""" |
| def __init__(self, o=128): |
| super().__init__() |
| c = [5, 32, 64, 96, 128] |
| layers = [] |
| for a, b in zip(c[:-1], c[1:]): |
| layers += [nn.Conv2d(a, b, 3, padding=1), nn.GELU(), nn.Conv2d(b, b, 3, padding=1, stride=2), nn.GELU()] |
| self.net = nn.Sequential(*layers) |
| self.proj = nn.Linear(128 * 2 * 2, o) |
|
|
| def forward(self, x): |
| B, M = x.shape[:2] |
| x = x.reshape(B * M, 4, 32, 32) |
| a = x[:, 3:4] |
| lum = (0.2126 * x[:, 0:1] + 0.7152 * x[:, 1:2] + 0.0722 * x[:, 2:3]) * a |
| x = torch.cat([x[:, :3] * a, a, lum], 1) |
| x = self.net(x) |
| return self.proj(x.flatten(1)).reshape(B, M, -1) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, d, h, ff, drop): |
| super().__init__() |
| self.h = h |
| self.ln1 = nn.LayerNorm(d) |
| self.qkv = nn.Linear(d, 3 * d) |
| self.proj = nn.Linear(d, d) |
| self.rel = nn.Linear(K_REL, h, bias=False) |
| nn.init.normal_(self.rel.weight, std=0.5) |
| self.ln2 = nn.LayerNorm(d) |
| self.ff = nn.Sequential(nn.Linear(d, ff * d), nn.GELU(), nn.Dropout(drop), nn.Linear(ff * d, d)) |
| self.drop = nn.Dropout(drop) |
|
|
| def forward(self, x, rel, pad_bias): |
| B, T, D = x.shape |
| q, k, v = self.qkv(self.ln1(x)).reshape(B, T, 3, self.h, D // self.h).permute(2, 0, 3, 1, 4) |
| bias = self.rel(rel).permute(0, 3, 1, 2) + pad_bias |
| a = F.scaled_dot_product_attention(q, k, v, attn_mask=bias.to(q.dtype)) |
| x = x + self.drop(self.proj(a.transpose(1, 2).reshape(B, T, D))) |
| return x + self.drop(self.ff(self.ln2(x))) |
|
|
|
|
| class NadirNet(nn.Module): |
| def __init__(self, d=384, layers=8, heads=6, ff=4, drop=0.1): |
| super().__init__() |
| self.cfg = dict(d=d, layers=layers, heads=heads, ff=ff, drop=drop) |
| self.tok_in = MLP(FT + N_ROLES + 36, 2 * d, d) |
| self.cnn = ThumbCNN(128) |
| self.img_in = MLP(FI + 128, 2 * d, d) |
| self.page_in = MLP(FP, d, d) |
| self.src = nn.Embedding(2, d) |
| self.typ = nn.Embedding(3, d) |
| self.blocks = nn.ModuleList([Block(d, heads, ff, drop) for _ in range(layers)]) |
| self.ln = nn.LayerNorm(d) |
| self.color_head = MLP(d + 3, d, 6) |
| self.ctx_head = nn.Linear(d, N_CTX) |
| self.elem_head = MLP(d + N_CTX + 72, d, N_TREAT + N_IROLE, drop=drop) |
|
|
| def embed_images(self, img_f, img_thumb): |
| return self.img_in(torch.cat([img_f, self.cnn(img_thumb)], -1)) |
|
|
| def elements(self, e_f, e_thumb, under_light, under_dark, ctx): |
| """e_f (B,K,FI), e_thumb (B,K,4,32,32), under_light/under_dark (B,K,3) OKLab, ctx (B,N_CTX) |
| -> treatment logits (B,K,N_TREAT), role logits (B,K,N_IROLE).""" |
| K = e_f.shape[1] |
| h = torch.cat([self.embed_images(e_f, e_thumb), ctx[:, None].expand(-1, K, -1), fourier_lab(under_light), fourier_lab(under_dark)], -1) |
| z = self.elem_head(h) |
| return z[..., :N_TREAT], z[..., N_TREAT:] |
|
|
| def forward(self, tok_f, tok_role, tok_lab, img_f, img_thumb, page_f, src, rel, mask): |
| """tok_f (B,N,FT) tok_role (B,N) long, tok_lab (B,N,3), img_f (B,M,FI), img_thumb (B,M,4,32,32) float 0-1, |
| page_f (B,FP), src (B,) long, rel (B,T,T,K) with T=1+N+M, mask (B,T) bool (True=real).""" |
| B, N = tok_role.shape |
| M = img_f.shape[1] |
| tok = self.tok_in(torch.cat([tok_f, F.one_hot(tok_role, N_ROLES).to(tok_f.dtype), fourier_lab(tok_lab)], -1)) + self.typ.weight[1] |
| if M > 0: |
| img = self.embed_images(img_f, img_thumb) + self.typ.weight[2] |
| else: |
| img = tok.new_zeros(B, 0, tok.shape[-1]) |
| cls = self.page_in(page_f)[:, None] + self.typ.weight[0] |
| x = torch.cat([cls, tok, img], 1) + self.src(src)[:, None] |
| pad_bias = torch.zeros(mask.shape, dtype=x.dtype, device=x.device).masked_fill(~mask, -1e4)[:, None, None, :] |
| for blk in self.blocks: |
| x = blk(x, rel, pad_bias) |
| x = self.ln(x) |
| ht = x[:, 1:1 + N] |
| z = self.color_head(torch.cat([ht, tok_lab], -1)) |
| L = torch.sigmoid(z[..., 0]) |
| s = 1.6 * torch.sigmoid(z[..., 1]) |
| r = 0.6 * torch.tanh(z[..., 2]) |
| a_in, b_in = tok_lab[..., 1], tok_lab[..., 2] |
| a = s * a_in - r * b_in + 0.12 * torch.tanh(z[..., 3]) |
| b = r * a_in + s * b_in + 0.12 * torch.tanh(z[..., 4]) |
| lab = torch.stack([L, a, b], -1) |
| alpha_logit = z[..., 5] |
| ctx = self.ctx_head(x[:, 0]) |
| return lab, alpha_logit, ctx |
|
|
|
|
| class ElemNet(nn.Module): |
| """The element head as its own graph (shares NadirNet's weights): classify any set of page images.""" |
| def __init__(self, net): |
| super().__init__() |
| self.net = net |
|
|
| def forward(self, e_f, e_thumb, under_light, under_dark, ctx): |
| return self.net.elements(e_f, e_thumb, under_light, under_dark, ctx) |
|
|
|
|
| |
| OKLAB_TO_LMS = torch.tensor([[1.0, 0.3963377774, 0.2158037573], [1.0, -0.1055613458, -0.0638541728], [1.0, -0.0894841775, -1.2914855480]]) |
| LMS_TO_RGB = torch.tensor([[4.0767416621, -3.3077115913, 0.2309699292], [-1.2684380046, 2.6097574011, -0.3413193965], [-0.0041960863, -0.7034186147, 1.7076147010]]) |
|
|
|
|
| def oklab_to_linrgb(lab): |
| lms = lab @ OKLAB_TO_LMS.to(lab).T |
| return (lms ** 3) @ LMS_TO_RGB.to(lab).T |
|
|
|
|
| def luminance(lab): |
| rgb = oklab_to_linrgb(lab).clamp(0, 1) |
| return rgb @ torch.tensor([0.2126, 0.7152, 0.0722]).to(lab) |
|
|
|
|
| def contrast_ratio(lab1, lab2): |
| l1, l2 = luminance(lab1), luminance(lab2) |
| return (torch.maximum(l1, l2) + 0.05) / (torch.minimum(l1, l2) + 0.05) |
|
|
