pixal3d_mrp_mlx/tests/test_proj.py

"""Numerical oracle for the proj conditioning — diffed against upstream, not transcribed.

The proj extractor contains NO sparse convolution, so upstream runs on CPU torch here
and is a real oracle (same situation as the flow models, unlike the sparse decoders).
CLAUDE.md's rule applies: do not reason about correctness, measure it. Three of the
bugs found in this port were invisible to weight-key matching and only fell out of a
correlation against upstream.

Each check reports max abs diff and correlation against the real thing.
"""

import sys
from pathlib import Path

import mlx.core as mx
import numpy as np
import torch
import torch.nn.functional as F

REPO = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO))
sys.path.insert(0, str(REPO / "upstream" / "Pixal3D"))

from pixal3d_mlx import proj as P  # noqa: E402
from pixal3d.trainers.flow_matching.mixins.image_conditioned_proj import (  # noqa: E402
    ProjGrid as UpstreamProjGrid,
    project_points_to_image_batch,
    sample_features,
)

FOV = 0.6911  # ~39.6 deg, a typical Pixal3D estimate
PASS, FAIL = [], []


def _report(name, mine, ref, tol=2e-5):
    a, b = np.asarray(mine, dtype=np.float64).ravel(), np.asarray(ref, dtype=np.float64).ravel()
    diff = float(np.abs(a - b).max())
    corr = 1.0 if (a.std() == 0 and b.std() == 0) else float(np.corrcoef(a, b)[0, 1])
    ok = diff <= tol and corr > 0.999999
    (PASS if ok else FAIL).append(name)
    print(f"  {'PASS' if ok else 'FAIL'}  {name:<38} max diff {diff:.3e}  corr {corr:.8f}")


def test_bilinear_sample():
    """grid_sample(align_corners=False, padding_mode='border').

    Deliberately includes coordinates well outside [-1,1] — that is the only way the
    border-padding and clamp-before-corners behaviour is exercised at all. A naive
    lerp passes an in-bounds-only test and fails silently on every silhouette edge,
    which is exactly where a pixel-aligned model lives.
    """
    rng = np.random.default_rng(0)
    b, h, w, c, k = 2, 9, 11, 5, 400
    fmap = rng.standard_normal((b, h, w, c)).astype(np.float32)
    grid = rng.uniform(-1.8, 1.8, (b, k, 2)).astype(np.float32)  # deliberately out of range

    mine = P.bilinear_sample(mx.array(fmap), mx.array(grid))
    ref = sample_features(
        torch.from_numpy(fmap).permute(0, 3, 1, 2), torch.from_numpy(grid)
    ).permute(0, 2, 1)
    _report("bilinear_sample vs grid_sample", mine, ref.numpy())

    # and again strictly inside, so an out-of-range-only success can't hide an
    # interior interpolation error
    grid_in = rng.uniform(-0.95, 0.95, (b, k, 2)).astype(np.float32)
    mine_in = P.bilinear_sample(mx.array(fmap), mx.array(grid_in))
    ref_in = sample_features(
        torch.from_numpy(fmap).permute(0, 3, 1, 2), torch.from_numpy(grid_in)
    ).permute(0, 2, 1)
    _report("bilinear_sample (interior only)", mine_in, ref_in.numpy())


def test_project_points():
    rng = np.random.default_rng(1)
    b, n, res = 2, 500, 512
    pts = rng.uniform(-1, 1, (b, n, 3)).astype(np.float32)
    tm = np.broadcast_to(P._FRONT_VIEW, (b, 4, 4)).copy()
    tm[:, 1, 3] = -1.7

    px, depth, valid = P.project_points(mx.array(pts), mx.array(tm), FOV, res)
    r_px, r_depth, r_valid = project_points_to_image_batch(
        torch.from_numpy(pts), torch.from_numpy(tm), torch.tensor([FOV] * b), res
    )
    _report("project_points pixels", px, r_px.numpy(), tol=1e-3)
    _report("project_points depth", depth, r_depth.numpy(), tol=1e-5)
    same = bool((np.asarray(valid) == r_valid.numpy()).all())
    (PASS if same else FAIL).append("valid mask")
    print(f"  {'PASS' if same else 'FAIL'}  {'project_points valid mask':<38} exact match={same}")


def test_distance_from_fov():
    """Upstream's helper lives in inference.py; replicate its exact call here."""
    rot = torch.tensor([[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]])
    res, mesh_scale = 512, 1.0
    gp = torch.tensor([-1.0, 0.0, 0.0]) @ rot.T
    gp = gp / mesh_scale / 2
    f_pixels = float((16.0 / torch.tan(torch.tensor(FOV / 2.0))) * res / 32.0)
    x_ndc = 0.0 - res / 2.0
    ref = f_pixels * float(gp[0]) / x_ndc - float(gp[1])
    mine = P.distance_from_fov(FOV, mesh_scale, res)
    _report("distance_from_fov", np.array([mine]), np.array([ref]), tol=1e-6)


def test_proj_grid_full():
    """The whole ProjGrid forward at the ss stage's real shapes."""
    rng = np.random.default_rng(2)
    b, res, grid_res, d = 1, 512, 16, 1024
    patch_n = res // 16
    fmap = rng.standard_normal((b, patch_n, patch_n, d)).astype(np.float32)
    distance = P.distance_from_fov(FOV, 1.0, res)

    mine = P.ProjGrid(grid_resolution=grid_res, image_resolution=res)(
        mx.array(fmap), FOV, distance, 1.0
    )
    up = UpstreamProjGrid(grid_resolution=grid_res, image_resolution=res)
    ref = up(torch.from_numpy(fmap), torch.tensor([FOV]),
             torch.tensor([distance]), torch.tensor([1.0]), None)
    _report("ProjGrid forward (ss, 16^3)", mine, ref.numpy())

    # the lattice itself must match before the projection can mean anything
    _report("grid lattice", P.ProjGrid(grid_res, res).grid_points, up.grid_points.numpy(), tol=1e-7)


def test_full_extractor():
    """End-to-end: our DINOv3 + ProjGrid vs upstream's DinoV3ProjFeatureExtractor.

    This is the check that matters — it covers token splitting, the parameterless
    final layer_norm, the patch reshape and the projection together, on the real
    ViT-L/16 weights.

    Upstream's `extract_features` reaches for `self.model.layer`, which transformers
    5.x moved to `self.model.model.layer`. We restore the attribute rather than
    rewriting the method, so upstream's own unmodified code is what executes.
    """
    from pixal3d.trainers.flow_matching.mixins.image_conditioned_proj import (
        DinoV3ProjFeatureExtractor,
    )

    from pixal3d_mlx.cond import ProjConditioner
    from pixal3d_mlx.dino import MODEL_NAME, _layers

    res, grid_res = 512, 16
    rng = np.random.default_rng(3)
    img = rng.uniform(0, 1, (1, 3, res, res)).astype(np.float32)
    distance = P.distance_from_fov(FOV, 1.0, res)

    up = DinoV3ProjFeatureExtractor(model_name=MODEL_NAME, image_size=res,
                                    grid_resolution=grid_res).eval()
    if not hasattr(up.model, "layer"):
        up.model.layer = _layers(up.model)   # transformers 5.x attribute move

    r_global, r_proj = up(torch.from_numpy(img), torch.tensor([FOV]),
                          torch.tensor([distance]), torch.tensor([1.0]), None)

    cond, uncond = ProjConditioner("ss", device="cpu")(img, FOV, distance, 1.0)

    _report("extractor global tokens", cond["global"], r_global.numpy(), tol=1e-4)
    _report("extractor proj features", cond["proj"], r_proj.numpy(), tol=1e-4)

    shapes_ok = (cond["global"].shape == tuple(r_global.shape)
                 and cond["proj"].shape == tuple(r_proj.shape)
                 and cond["proj"].shape[1] == grid_res ** 3)
    (PASS if shapes_ok else FAIL).append("extractor shapes")
    print(f"  {'PASS' if shapes_ok else 'FAIL'}  {'extractor shapes':<38} "
          f"global {cond['global'].shape}  proj {cond['proj'].shape}")

    zeroed = float(np.abs(np.asarray(uncond["proj"])).max()) == 0.0
    (PASS if zeroed else FAIL).append("uncond is zeros")
    print(f"  {'PASS' if zeroed else 'FAIL'}  {'uncond is zeros':<38} {zeroed}")


def test_naf_stage_wiring():
    """A NAF stage (shape_512) end to end against upstream's extractor.

    SCOPE: `_patch_naf_attention` replaces NAF's natten call globally, so upstream's
    extractor runs through the SAME neighborhood kernel we do — this test therefore
    verifies the WIRING (guide resizing, which proj_grid the high-res branch reuses,
    concat order, proj_channels = embed_dim*2), NOT the kernel. The kernel is verified
    separately against natten itself in tests/test_naf.py, which is where that
    equivalence has to be established.
    """
    from pixal3d.trainers.flow_matching.mixins.image_conditioned_proj import (
        DinoV3ProjFeatureExtractor,
    )

    from pixal3d_mlx.cond import CONFIGS, ProjConditioner
    from pixal3d_mlx.dino import MODEL_NAME, _layers
    from pixal3d_mlx.naf import _patch_naf_attention, load_naf

    cfg = CONFIGS["shape_512"]
    res, grid_res = cfg["image_size"], cfg["grid_resolution"]
    rng = np.random.default_rng(5)
    img = rng.uniform(0, 1, (1, 3, res, res)).astype(np.float32)
    distance = P.distance_from_fov(FOV, 1.0, res)

    load_naf("cpu")
    _patch_naf_attention()

    up = DinoV3ProjFeatureExtractor(
        model_name=MODEL_NAME, image_size=res, grid_resolution=grid_res,
        use_naf_upsample=True, naf_target_size=cfg["naf_target_size"],
    ).eval()
    if not hasattr(up.model, "layer"):
        up.model.layer = _layers(up.model)
    up._load_naf()

    r_global, r_proj = up(torch.from_numpy(img), torch.tensor([FOV]),
                          torch.tensor([distance]), torch.tensor([1.0]), None)

    cond, _ = ProjConditioner("shape_512", device="cpu")(img, FOV, distance, 1.0)

    _report("naf stage global", cond["global"], r_global.numpy(), tol=1e-4)
    _report("naf stage proj (lr+hr concat)", cond["proj"], r_proj.numpy(), tol=1e-3)

    d = up.embed_dim
    ok = (cond["proj"].shape == (1, grid_res ** 3, d * 2) == tuple(r_proj.shape))
    (PASS if ok else FAIL).append("naf proj_channels")
    print(f"  {'PASS' if ok else 'FAIL'}  {'naf proj_channels = embed_dim*2':<38} "
          f"{cond['proj'].shape}")


if __name__ == "__main__":
    torch.set_grad_enabled(False)
    for fn in (test_bilinear_sample, test_project_points,
               test_distance_from_fov, test_proj_grid_full, test_full_extractor,
               test_naf_stage_wiring):
        fn()
    print(f"\n{len(PASS)}/{len(PASS) + len(FAIL)} passed")
    if FAIL:
        print("FAILED: " + ", ".join(FAIL))
        sys.exit(1)