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.
30 lines
1.6 KiB
Markdown
30 lines
1.6 KiB
Markdown
# 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.py` is 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.
|
|
- `_scale` uses `Fraction`, 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/Pixal3D` is vendored read-only reference; never edit it.
|
|
- Weights live at MODELBEAST/vendor/pixal3d-weights (24GB, outside this repo).
|