6f40fae0cc
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>
99 lines
3.6 KiB
Python
99 lines
3.6 KiB
Python
"""The low-res reduction that replaces natten, diffed against natten itself.
|
|
|
|
natten cannot run this model on Apple Silicon (no libnatten -> cutlass-fna raises;
|
|
flex-fna is CPU-only, rejects NAF's asymmetric head dims, and was OOM-killed at 512).
|
|
So `lr_neighborhood_attention` replaces it. That is only defensible if it is EXACTLY
|
|
equivalent, which is what this asserts — against natten's own kernel, at the sizes
|
|
where natten does run.
|
|
|
|
Skips cleanly if natten is absent; the reduction does not need it at runtime.
|
|
"""
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
REPO = Path(__file__).resolve().parents[1]
|
|
sys.path.insert(0, str(REPO))
|
|
|
|
from pixal3d_mlx.naf import lr_neighborhood_attention # noqa: E402
|
|
|
|
PASS, FAIL = [], []
|
|
|
|
|
|
def _natten_reference(q, k_lr, v_lr, n, ks):
|
|
"""What NAF actually does: upsample K/V, then dilated neighborhood attention."""
|
|
import natten
|
|
import torch.nn.functional as F
|
|
|
|
hq, wq = q.shape[-2:]
|
|
hk, wk = k_lr.shape[-2:]
|
|
|
|
def to_bhwnd(x):
|
|
return x.reshape(x.shape[0], n, -1, x.shape[-2], x.shape[-1]).permute(0, 3, 4, 1, 2)
|
|
|
|
kk = to_bhwnd(F.interpolate(k_lr, size=(hq, wq), mode="nearest-exact"))
|
|
vv = to_bhwnd(F.interpolate(v_lr, size=(hq, wq), mode="nearest-exact"))
|
|
out = natten.na2d(to_bhwnd(q), kk, vv, kernel_size=ks,
|
|
dilation=(hq // hk, wq // wk), stride=1, backend="flex-fna")
|
|
b = out.shape[0]
|
|
return out.permute(0, 3, 4, 1, 2).reshape(b, -1, hq, wq)
|
|
|
|
|
|
def test_matches_natten():
|
|
"""natten requires kernel*dilation <= input, so low-res must be >= kernel."""
|
|
try:
|
|
import natten # noqa: F401
|
|
except ImportError:
|
|
print(" SKIP natten not installed — reduction is self-sufficient at runtime")
|
|
return
|
|
|
|
torch.manual_seed(0)
|
|
n, ks, dq = 4, (9, 9), 16
|
|
for hk, hq in ((16, 64), (16, 128), (32, 256)):
|
|
q = torch.randn(1, n * dq, hq, hq)
|
|
k = torch.randn(1, n * dq, hk, hk)
|
|
v = torch.randn(1, n * dq, hk, hk)
|
|
mine = lr_neighborhood_attention(q, k, v, n, ks)
|
|
ref = _natten_reference(q, k, v, n, ks)
|
|
diff = float((mine - ref).abs().max())
|
|
ok = diff < 5e-6
|
|
(PASS if ok else FAIL).append(f"natten dil {hq // hk}")
|
|
print(f" {'PASS' if ok else 'FAIL'} vs natten LR {hk:3} -> HR {hq:4} "
|
|
f"(dil {hq // hk:2}) max diff {diff:.3e}")
|
|
|
|
|
|
def test_asymmetric_head_dims():
|
|
"""NAF uses qk=64 with v=256 — the exact case flex-fna refuses."""
|
|
torch.manual_seed(1)
|
|
n, hk, hq = 4, 16, 64
|
|
q = torch.randn(1, n * 64, hq, hq)
|
|
k = torch.randn(1, n * 64, hk, hk)
|
|
v = torch.randn(1, n * 256, hk, hk)
|
|
out = lr_neighborhood_attention(q, k, v, n, (9, 9))
|
|
ok = out.shape == (1, n * 256, hq, hq) and bool(torch.isfinite(out).all())
|
|
(PASS if ok else FAIL).append("asymmetric head dims")
|
|
print(f" {'PASS' if ok else 'FAIL'} asymmetric qk=64 v=256 -> {tuple(out.shape)}")
|
|
|
|
|
|
def test_rejects_non_integer_ratio():
|
|
"""The reduction is only valid when the dilation is exactly the upsample factor."""
|
|
q = torch.randn(1, 4 * 16, 50, 50)
|
|
k = torch.randn(1, 4 * 16, 16, 16)
|
|
try:
|
|
lr_neighborhood_attention(q, k, k, 4, (9, 9))
|
|
ok = False
|
|
except ValueError:
|
|
ok = True
|
|
(PASS if ok else FAIL).append("guards non-integer ratio")
|
|
print(f" {'PASS' if ok else 'FAIL'} rejects non-integer target ratio")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
torch.set_grad_enabled(False)
|
|
for fn in (test_matches_natten, test_asymmetric_head_dims, test_rejects_non_integer_ratio):
|
|
fn()
|
|
print(f"\n{len(PASS)}/{len(PASS) + len(FAIL)} passed")
|
|
sys.exit(1 if FAIL else 0)
|