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.
1.6 KiB
1.6 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).