Completes the conditioning. shape_512 / shape_1024 / tex_1024 run a second HIGH-RES
branch — NAF upsamples the DINOv3 patch map to 512/1024 guided by the RGB image, the
proj grid samples that too, and the branches concatenate. That is why those stages
have proj_channels = embed_dim*2 (2048).
CORRECTS AN EARLIER CLAIM: I said natten is never imported and can be skipped. True of
Pixal3D's own source — but natten is a dependency of NAF (valeoai/NAF), which arrives
at RUNTIME via torch.hub and is not vendored. That is what README Step 3 is for. The
warning was real; the reason was one level down.
natten is a dead end on Apple Silicon regardless:
cutlass-fna requires libnatten, which the arm64 build does not produce
flex-fna CPU only ('not on a CUDA, ROCm, or CPU device: mps'), AND refuses
different head dims for QK vs V — which is exactly NAF's shape
(qk=64, v=256). Worked around, CPU took 243s at 256px and was
OOM-KILLED (exit 137) at the 512 the pipeline actually needs.
REPLACED BY AN EXACT REDUCTION, not an approximation. NAF resizes K/V from the 32x32
patch map with nearest-exact and then dilates by exactly the upsample factor, so the
dilated high-res neighborhood samples one position per low-res cell and collapses to a
plain clamped 9x9 neighborhood on the 32x32 grid, shared by every high-res pixel in
that cell. Verified against natten's own kernel at three dilations:
LR 16 -> HR 64 (dil 4) max diff 7.153e-07
LR 16 -> HR 128 (dil 8) max diff 7.153e-07
LR 32 -> HR 256 (dil 8) max diff 7.153e-07 (float32 epsilon)
Grouping queries by low-res cell also avoids materialising the high-res neighborhood,
which would be ~87GB of gathered V at 512. Result, on MPS:
32 -> 512 2.4s (natten: OOM-killed)
64 -> 512 0.2s
64 -> 1024 0.9s
Also fixes a caching bug that stranded NAF on whichever device loaded first, and one
in my own wiring: the high-res branch must reuse the SAME ProjGrid at image_size, not
a new one at naf_target_size. The normalised coordinate carries a 1/resolution term,
so the latter lands ~0.001 off in [-1,1] — a sub-pixel shift on every voxel, in the
one model whose entire premise is pixel alignment.
27/27 green (15 proj incl. the NAF stage vs upstream at corr 1.00000000, 5 sampler,
5 naf vs natten, 2 decoders).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
6f40fae0cc