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>
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>
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>
The structure stage now runs end to end on Metal: noise -> ss_flow (12-step Euler)
-> latent -> ss_dec -> 64^3 occupancy grid. 759 ms/step for the 1.3B DiT, 50 ms for
the decoder. Output lands at 1.36% occupancy, which is the right order for a surface
in a 64^3 grid.
Sampler details worth recording, both from upstream:
- CFG is a LERP (g*pos + (1-g)*neg), NOT neg + g*(pos-neg). Those differ non-linearly
in strength rather than failing outright, so it is silent when wrong. Tested.
- Guidance interval forces strength to 1 outside its window, which halves model calls
there: 6 calls for 4 steps rather than 8. Tested by counting.
Schedule matches upstream to 1e-16 and constant velocity integrates exactly at any
step count.
Remaining for a real image->3D run is CONDITIONING, not models: DINOv3 features plus
Pixal3D's camera back-projection for the view-aligned 'proj' half. Deliberately not
porting DINOv3 - it is a stock ViT run once per image, outside the denoising loop, so
torch on MPS is the right tool and transformers gives exact parity for free.
All seven Pixal3D models now load and run. shape_dec 292/292, tex_dec 284/284, both
with zero missing/unmapped/mismatched keys - the 8-param difference between them is
exactly the four to_subdiv layers, since tex_dec has pred_subdiv=False.
Behaviour is right: 12 latent voxels grow SELECTIVELY through four stages
(12 -> 91 -> 193 -> 358 -> 1150) rather than x8 each time, which would have reached
49,152. Scale lands at 1/16. shape_dec emits the hardcoded 7 channels and its vertex
head produces offsets inside the [-0.5, 1.5] band its sigmoid+voxel_margin allows.
tex_dec, guided by shape_dec's masks, reproduces exactly the same voxel count.
VERIFICATION CAVEAT, recorded in the README: these two are the only models using sparse
conv, so upstream cannot run here and there is no numerical oracle. Unlike the five
models verified at correlation 1.0, these are checked structurally and behaviourally
only. Weaker evidence, and labelled as such rather than presented alongside the
verified results.
slat_flow joins ss_flow: max abs diff 9.3e-6, 700/700 params, against upstream running
on CPU torch. All three SLAT checkpoints load clean and run (img2shape 512/1024 and
imgshape2tex, the last taking 64 in-channels since shape is concatenated).
The SLAT flows differ from ss_flow in two ways, both handled in the shared core:
tokens are a SparseTensor's voxels so attention runs per batch item, and RoPE phases
are NOT shipped - positions are the input's own coordinates, so they are derived at
call time by rope_phases_from_coords.
tests/oracle_slat.py keeps the CPU-torch oracle harness: it patches both the dense and
sparse flash-attn kernels with SDPA equivalents. Use it rather than reasoning about
correctness - it has already caught three bugs a perfect 700/700 key match did not.
