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 .json with 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