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.
2.2 KiB
2.2 KiB
pixal3d_mrp_mlx — working notes
MLX port of TencentARC Pixal3D, on top of the shared trellis_sparse_mlx core.
Measured facts (don't re-derive)
- Pixal3D inherits TRELLIS.2's sparse module, same as LATO.2.
conv_spconv.pyis the same code:stride==1 and padding is None -> spconv.SubMConv3d. - EVERY SparseConv3d in the model is
(channels, out, 3)-> submanifold only. The shared core's SubMConv3d covers all of them. - attn_mode is 'full' everywhere -> SDPA. No flash-attn/xformers needed.
- Container differs from LATO.2:
SparseTensor(VarLenTensor). VarLenTensor is feats + explicit slice-per-batch layout, no coords. Both are in the shared core now. - Cache API spelling is
get_spatial_cache/register_spatial_cache(LATO.2 uses cache_get/cache_put). Shared core exposes both. _scaleusesFraction, not int, upstream. Shared core keeps scale as an opaque tuple so either works.- Only sparse-conv-using model file is
models/sc_vaes/sparse_unet_vae.py.
Remaining ops to port
SparseUpsample(2)— pixal3d/modules/sparse/spatial/basic.py. Cache-paired: needs the subdivision stored by its matching SparseDownsample, or an explicit subdivision tensor.SparseSpatial2Channel(2)— pixal3d/modules/sparse/spatial/spatial2channel.py. Sparse pixel-shuffle; also cache-backed.
Conventions
- Depends on ../trellis_sparse_mlx (editable install). Do NOT vendor a second copy of the sparse core.
upstream/Pixal3Dis vendored read-only reference; never edit it.- Weights live at MODELBEAST/vendor/pixal3d-weights (24GB, outside this repo).
Weights (measured 2026-08-02)
- Every checkpoint ships a sibling .json with {name, args} — the authoritative config. Do NOT infer architecture from shapes here (that was needed for LATO.2, not this).
- ss_flow + slat_flow x3 (~20GB, the bulk): ZERO 5-D tensors. Pure transformers, no sparse conv. They pass through the converter untouched.
- shape_dec / tex_dec: 40 KRSC sparse kernels each. ss_dec: 20 dense Conv3d. classify_5d() distinguishes them correctly on the real weights.
rope_phasesships as a stored tensor — RoPE phases are precomputed, not derived.- Converter preserves dtype (fp16 stays fp16); upcasting doubled 24GB for nothing.