Commit Graph

3 Commits

Author SHA1 Message Date
m3ultra
6f40fae0cc All four proj extractors: NAF high-res branch, without natten
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>
2026-08-03 14:39:29 +10:00
m3ultra
79ef81988c Real image -> occupancy grid, with a silhouette check and honest timings
image_to_occupancy() runs the structure stage on an actual photo: preprocess ->
DINOv3 -> proj back-projection -> ss_flow -> ss_dec -> 64^3 occupancy.

VERIFICATION THAT MATTERS: scripts/run_structure.py re-projects the occupied voxels
through the same camera and compares against the input alpha matte. On the upstream
sample that is silhouette IoU 0.842 with 12948 voxels occupied (4.94% of 64^3). This
is the model's own headline claim, so it is the right thing to assert — 'it ran
without crashing' would pass just as happily on a generic blob.

Two real bugs this phase found, neither visible without reading the shipped configs:

1. THE SAMPLER WAS MISSING guidance_rescale. The checkpoint's own pipeline.json sets
   0.7 for the structure stage and 0.5 for shape_slat, so this fires at the model's
   DEFAULT settings — omitting it silently overcooks every structure prediction. Now
   implemented (Lin et al. CFG rescale) and diffed against upstream's
   ClassifierFreeGuidanceSamplerMixin, run directly rather than reimplemented.
2. The sampler defaults were wrong: the real ss stage is steps=12 / rescale_t=5.0 /
   guidance 7.5 / interval [0.6,1.0], not the steps=25 / rescale_t=3.0 the smoke test
   assumed. All three stages' real params now live in pipeline.py, read from
   pipeline.json rather than guessed.

TIMINGS, measured with interleaved reps after warmup (the first pass attributed the
same 11s of residual warmup to both 'rescale' and 'torch contention'; it was neither):

  cold run                    89.3s
  warm, full settings         16.5s
  warm, CFG off                9.2s   -> CFG costs 1.80x, as expected for 10/12
                                         steps falling inside the guidance interval
  guidance_rescale              ~0s   -> free
  torch/MPS contention          ~0s   -> DINOv3 can stay resident
  peak memory                  6.8GB

THE FINDING THAT SHAPES THE OPERATOR: warmup is ~71s against ~17s of actual compute,
i.e. 4x the work. A MODELBEAST operator MUST hold the models resident across jobs
rather than fork per job — the trellis2 lane shows the same shape (47.9s cold vs 2.5s
warm pipeline_load). Cost this in before optimising any kernel.

17/17 tests green (12 proj + 5 sampler).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-08-03 14:10:30 +10:00
m3ultra
bd0f3f1941 Proj conditioning: camera back-projection, grid_sample, DINOv3 (12/12 vs upstream)
The pixel-aligned conditioning is the ONLY thing separating this port from the
trellis2_mlx operator already in MODELBEAST — upstream's main branch is the
TRELLIS.2 backbone, so everything else here is TRELLIS.2 with a different head.
This lands that head.

  proj.py  ProjGrid, project_points, bilinear_sample, distance_from_fov — MLX
  dino.py  DINOv3 ViT-L/16 left in torch on MPS (run once per image, outside the
           25-step loop; transformers gives exact parity for free)
  cond.py  encode_image_proj equivalent -> {'global','proj'} + zero uncond

The extractor has no sparse conv, so upstream RUNS on CPU torch here and is a real
oracle. All 12 checks diff against it, not against a transcription:

  bilinear_sample vs grid_sample     max diff 2.4e-07   corr 1.00000000
  project_points pixels/depth/mask   exact
  ProjGrid forward (ss, 16^3)        max diff 1.9e-05   corr 1.00000000
  extractor global tokens            max diff 0.0e+00   corr 1.00000000
  extractor proj features            max diff 4.8e-06   corr 1.00000000

Three details that a plain transcription gets wrong and eyeballing cannot catch:
grid_sample's align_corners=False maps a normalised coord to ((c+1)*size-1)/2, not
(c+1)/2*(size-1) — half a texel, invisible until you compare; padding_mode='border'
clamps the SOURCE INDEX before corners are taken, not the corners after, which
changes the weights on every silhouette edge (tested with deliberately out-of-range
grid coords); and the camera looks down -Z, so a sign slip still yields a plausible
grid that samples the mirror image.

Also corrects a shape assumption from the earlier smoke test: 'global' is CLS + 4
register tokens = [B,5,1024], NOT the 1370 image tokens. The patch tokens go to the
proj branch. That asymmetry IS the architecture.

Note the parameterless final layer_norm in extract_features — not model.norm, which
has weights. Same trap as the ss_flow bug: no checkpoint trace, 200x output error.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-08-03 13:52:41 +10:00