"""NAF feature upsampler — the reason Pixal3D's README tells you to install natten. The three high-resolution extractor stages (shape_512, shape_1024, tex_1024) run a second, HIGH-RES branch: NAF upsamples the 32x32 DINOv3 patch map to 512 or 1024 using the RGB image as a guide, the proj grid samples that too, and the two branches are concatenated — which is why those stages have `proj_channels = embed_dim * 2`. Worth being precise about, because it is easy to get the opposite impression from the repo: `natten` is NEVER imported by any file in Pixal3D itself. It is a dependency of NAF (valeoai/NAF), which arrives at runtime through `torch.hub`, not vendored. That is what README Step 3's `NATTEN_CUDA_ARCH=... pip install natten` is really for. natten IS A DEAD END ON THIS MACHINE, so it is not used at all: * NAF hardcodes `backend="cutlass-fna"`, a CUDA kernel. natten pip-installs and builds on arm64 but WITHOUT `libnatten`, so that backend raises. * Its `flex-fna` backend is pure PyTorch and does run — but only on CPU ("Can't run Flex Attention; tensor is not on a CUDA, ROCm, or CPU device: mps"), and it refuses different head dims for QK vs V, which NAF uses (qk=64, v=256). Even working around the head dims, CPU flex-fna took 243s at 256px and was OOM-killed at 512. WHAT REPLACES IT — an exact algebraic reduction, not an approximation: NAF resizes K and V from the 32x32 patch map up to the target with `nearest-exact`, then runs neighborhood attention with `dilation = target // 32` — i.e. the dilation is exactly the upsample factor. So the dilated high-res neighborhood samples exactly one position per low-res cell, and the whole thing collapses to a plain **clamped 9x9 neighborhood on the 32x32 grid**, shared by every high-res pixel in the same cell. Verified against natten's own flex-fna at three dilations (4, 8, 8) — max abs diff 7.15e-07, i.e. float32 epsilon. That equivalence is what makes this legitimate rather than a lookalike, and it is asserted in tests/test_naf.py. The reduction also makes it cheap: attention is computed once per (low-res cell, sub-pixel) pair instead of once per high-res pixel against a materialised high-res neighbourhood — which would be ~87GB of gathered V at 512. Runs on MPS. Kept in torch for the same reason as DINOv3: once per image, outside the denoising loop, so exact parity is worth more than a port. """ from __future__ import annotations import mlx.core as mx import numpy as np _naf_model = None def lr_neighborhood_attention(q, k_lr, v_lr, num_heads, kernel_size): """Clamped k x k neighborhood attention on the low-res grid. q [B, n*dq, H, W] high-res queries k_lr [B, n*dk, h, w] low-res keys (NOT pre-upsampled) v_lr [B, n*dv, h, w] low-res values -> [B, n*dv, H, W] Exactly equivalent to NAF's `na2d(q, upsample(k), upsample(v), dilation=H//h)`; see the module docstring. Requires H % h == 0, which holds for every shipped configuration (512/32, 512/64, 1024/64). """ import torch b, _, hq, wq = q.shape _, _, hk, wk = k_lr.shape kh, kw = kernel_size n = num_heads if hq % hk or wq % wk: raise ValueError(f"target {(hq, wq)} is not an integer multiple of {(hk, wk)}") rh, rw = hq // hk, wq // wk def split(x): return x.reshape(b, n, -1, x.shape[-2], x.shape[-1]).permute(0, 3, 4, 1, 2) qq, kk, vv = split(q), split(k_lr), split(v_lr) # [B,H,W,n,d] d = qq.shape[-1] # clamped window starts, so every cell attends to exactly kh*kw positions sh = (torch.arange(hk, device=q.device) - kh // 2).clamp(0, max(hk - kh, 0)) sw = (torch.arange(wk, device=q.device) - kw // 2).clamp(0, max(wk - kw, 0)) oh = sh[:, None] + torch.arange(kh, device=q.device)[None, :] ow = sw[:, None] + torch.arange(kw, device=q.device)[None, :] idx = (oh[:, None, :, None] * wk + ow[None, :, None, :]).reshape(-1) kn = kk.reshape(b, hk * wk, n, -1)[:, idx].reshape(b, hk, wk, kh * kw, n, -1) vn = vv.reshape(b, hk * wk, n, -1)[:, idx].reshape(b, hk, wk, kh * kw, n, -1) # group high-res queries by their low-res cell so V is never expanded to high-res qg = (qq.reshape(b, hk, rh, wk, rw, n, d) .permute(0, 1, 3, 2, 4, 5, 6) .reshape(b, hk, wk, rh * rw, n, d)) scores = torch.einsum("bacpnd,backnd->bacpnk", qg, kn) * (d ** -0.5) weights = scores.softmax(-1) out = torch.einsum("bacpnk,backnd->bacpnd", weights, vn) return (out.reshape(b, hk, wk, rh, rw, n, -1) .permute(0, 1, 3, 2, 4, 5, 6) .reshape(b, hq, wq, n, -1) .permute(0, 3, 4, 1, 2) .reshape(b, -1, hq, wq)) def _patch_naf_attention(): """Route NAF's CrossAttention through the low-res reduction instead of natten.""" from src.layers.attentions import CrossAttention if getattr(CrossAttention, "_pixal3d_patched", False): return def forward(self, q, k, v, image=None, return_weights=False, **kwargs): if return_weights: raise NotImplementedError("attention weights are not needed for inference") return lr_neighborhood_attention(q, k, v, self.num_heads, self.kernel_size) CrossAttention.forward = forward CrossAttention._pixal3d_patched = True def load_naf(device: str = "mps"): """Load NAF once per process (2.5MB checkpoint, pulled by torch.hub). The cached model is moved when a different device is asked for — caching it on whichever device happened to load first silently strands the weights and produces a device-mismatch deep inside the first conv. """ global _naf_model if _naf_model is None: import sys import torch from pathlib import Path hub = Path.home() / ".cache" / "torch" / "hub" / "valeoai_NAF_main" if hub.exists() and str(hub) not in sys.path: sys.path.insert(0, str(hub)) # NAF's modules import as `src.*` model = torch.hub.load("valeoai/NAF", "naf", pretrained=True, device=device, trust_repo=True).eval() model.requires_grad_(False) _patch_naf_attention() _naf_model = model if str(next(_naf_model.parameters()).device).split(":")[0] != device: _naf_model = _naf_model.to(device) return _naf_model def upsample(patch_bhwc: mx.array, image_chw01: np.ndarray, target: int, device: str = "mps") -> mx.array: """DINOv3 patch map [B,h,w,D] + guide image [B,3,H,W] -> [B,T,T,D] at target size. Returned channels-last so it drops straight into `ProjGrid`, which is BHWC. """ import torch model = load_naf(device) guide = torch.from_numpy(np.asarray(image_chw01, dtype=np.float32)).to(device) lr = torch.from_numpy(np.asarray(patch_bhwc, dtype=np.float32)).permute(0, 3, 1, 2).to(device) with torch.no_grad(): hr = model(guide, lr, (target, target)) # [B, D, T, T] return mx.array(hr.permute(0, 2, 3, 1).float().cpu().numpy())