pixal3d_mrp_mlx/pixal3d_mlx/naf.py

"""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())