pixal3d_mrp_mlx/CLAUDE.md
John 80991a8fe9 Scope Pixal3D against the shared sparse core
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.
2026-08-02 10:36:35 +10:00

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