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.
74/74 params, max abs diff 4.6e-4 on values around -163 (~3e-6 relative), 60ms for a
16^3 latent -> 64^3 occupancy grid.
Fully dense, so no sparse ops involved - but MLX's Conv3d is channels-LAST where torch
is NCDHW, so tensors are carried channels-last throughout and transposed only at the
boundaries. The converter already emits [O,kz,ky,kx,I] to match. pixel_shuffle_3d had
to be rewritten for that layout: the (H,s)(W,s)(D,s) interleave order is what matters
and getting it wrong scrambles the grid while preserving its shape.
This completes the structure stage end to end: image -> ss_flow -> latent -> ss_dec.
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.
700/700 params, max abs diff 1.2e-5 on the real 1.3B checkpoint.
Key realisation: the flow models have no sparse conv, so upstream RUNS on CPU torch
with flash-attn swapped for SDPA. That gives a real numerical oracle - unavailable for
the sparse path, where spconv cannot be installed at all.
It was needed. Three bugs survived a loader reporting a perfect 700/700 with zero
missing and zero unmapped keys:
- a parameterless final LayerNorm (no params -> no checkpoint trace) that the output
was 200x too large without
- rope_phases being complex64, so the rotation is a complex multiply
- qk_rms_norm belonging before rope rather than after
Weight-key matching is necessary but nowhere near sufficient for a port.
Converter is light because Pixal3D ships safetensors with sibling .json configs, so
there is no architecture to infer and no pickle to unpack. Only rank-5 tensors are
rearranged; dtype is preserved (upcasting fp16->fp32 doubled 24GB for no benefit) and
the KRSC remap is verified a pure permutation of values.
Measured: the four flow models (~20GB) have ZERO 5-D tensors - pure transformers that
never touch sparse conv. shape_dec/tex_dec carry 40 KRSC kernels each, ss_dec 20 dense
Conv3d, and classify_5d separates them correctly on the real weights.
Corrects the earlier 'gap is two ops' claim: that was right about modules/sparse but
undercounted the model blocks. The configs show all four flow models need RoPE,
qk_rms_norm and AdaLN modulation, and the decoders need SparseConvNeXtBlock3d and
SparseResBlockC2S3d.
Every SparseConv3d in the Pixal3D model is (channels, out, 3) - stride=1/padding=None,
so spconv dispatches to SubMConv3d, which trellis_sparse_mlx already implements and
benchmarks. attn_mode is 'full' throughout. The container split (VarLenTensor /
SparseTensor) and the get/register_spatial_cache spelling are already in the shared core.
Remaining gap is two ops: SparseUpsample and SparseSpatial2Channel, both cache-paired
with a matching downsample rather than recomputing structure.
