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

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