e23233731b
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.
3.5 KiB
pixal3d_mrp_mlx
MLX port of Pixal3D (TencentARC + Tsinghua, SIGGRAPH 2026, MIT) — pixel-aligned single-image 3D generation — for Apple Silicon.
Built on trellis_sparse_mlx, the shared TRELLIS-lineage
sparse core. Pixal3D and LATO.2 inherit the same sparse module from TRELLIS.2, so the
expensive part — submanifold sparse convolution, which has no Metal implementation — is
already done and tested there.
Why Pixal3D
It back-projects pixel features directly into 3D rather than injecting them through attention, so silhouettes stay exact to the source image. Different failure mode from TRELLIS/Hunyuan, and complementary to them.
Upstream needs ~24 GB VRAM, which is a wall on consumer Nvidia and a non-issue on a 128 GB+ Ultra.
Scope, measured against the shared core
| Need | Status |
|---|---|
SparseConv3d — every call is (c, out, 3), i.e. stride=1/padding=None → SubMConv3d |
✅ in shared core |
attn_mode='full' (the only mode used) |
✅ in shared core |
SparseLinear, norms, activations, ResBlock, transformer blocks |
✅ in shared core |
SparseDownsample(2) |
✅ in shared core |
VarLenTensor / SparseTensor split + get/register_spatial_cache |
✅ added to shared core |
dense nn.Conv3d(.., 2, stride=2) in sparse_structure_vae |
✅ maps to mlx.nn.Conv3d |
SparseUpsample(2) |
❌ to do — cache-paired inverse of a downsample |
SparseSpatial2Channel(2) |
❌ to do — sparse pixel-shuffle, spatial→channel |
Both of those are now done in the shared core.
Correction to the earlier scope
"the gap is two ops" was accurate about modules/sparse/ — the sparse primitives. It
undercounted the model blocks, which the configs revealed:
| Still needed | Where |
|---|---|
| RoPE positional embedding | all 4 flow models (pe_mode: "rope") — but rope_phases ships as a stored tensor, so phases are precomputed, not derived |
qk_rms_norm on q and k |
all 4 flow models |
AdaLN modulation (share_mod: true) |
all 4 flow models |
image_attn_mode: "proj" conditioning |
all 4 flow models |
SparseConvNeXtBlock3d |
shape_dec, tex_dec |
SparseResBlockC2S3d (channel↔spatial) |
shape_dec, tex_dec — uses SparseSpatial2Channel |
Offsetting that, a genuine simplification the tensors revealed: the four flow models
(~20 GB, the bulk of the download) contain ZERO 5-D tensors. ss_flow and the three
slat_flow DiTs are pure transformers — they never touch sparse convolution, so they
need none of the sparse core, just DiT blocks.
Model surface
pixal3d/models/
sparse_structure_vae.py dense Conv3d — voxel structure
sparse_structure_flow.py structure flow (SS)
structured_latent_flow.py SLAT flow
sc_vaes/sparse_unet_vae.py the only file using sparse conv
Weights: 24.04 GB across 19 files (1.3B DiTs at 512/1024 + shape/tex decoders).
Status
- Scoped against the shared core
- Weights downloaded (24 GB) — each ships a sibling
.jsonwith the exact config, so unlike LATO.2 there is no architecture to infer upsample(masked) +downsample(mode=)landed in the shared core- Weight converter — decoders remap KRSC→
[K³,in,out], dtype preserved, remap verified a pure permutation. Flow models pass through untouched. - DiT blocks: RoPE (stored phases), qk_rms_norm, AdaLN modulation
SparseConvNeXtBlock3d,SparseResBlockC2S3d,SparseSpatial2Channel- Model graphs
- End-to-end