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.
This commit is contained in:
John 2026-08-02 10:36:35 +10:00
commit 80991a8fe9
3 changed files with 91 additions and 0 deletions

6
.gitignore vendored Normal file
View File

@ -0,0 +1,6 @@
.venv
upstream/
ckpt/
weights/
__pycache__/
*.pyc

29
CLAUDE.md Normal file
View File

@ -0,0 +1,29 @@
# 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).

56
README.md Normal file
View File

@ -0,0 +1,56 @@
# pixal3d_mrp_mlx
MLX port of [Pixal3D](https://github.com/TencentARC/Pixal3D) (TencentARC + Tsinghua,
SIGGRAPH 2026, MIT) — pixel-aligned single-image 3D generation — for Apple Silicon.
Built on **[`trellis_sparse_mlx`](../trellis_sparse_mlx)**, the shared TRELLIS-lineage
sparse core. Pixal3D and LATO.2 inherit the same sparse module from TRELLIS.2, so the
expensive part — submanifold sparse convolution, which has no Metal implementation — is
already done and tested there.
## Why Pixal3D
It back-projects pixel features directly into 3D rather than injecting them through
attention, so silhouettes stay exact to the source image. Different failure mode from
TRELLIS/Hunyuan, and complementary to them.
Upstream needs ~24 GB VRAM, which is a wall on consumer Nvidia and a non-issue on a
128 GB+ Ultra.
## Scope, measured against the shared core
| Need | Status |
|---|---|
| `SparseConv3d`**every call is `(c, out, 3)`**, i.e. stride=1/padding=None → SubMConv3d | ✅ in shared core |
| `attn_mode='full'` (the only mode used) | ✅ in shared core |
| `SparseLinear`, norms, activations, ResBlock, transformer blocks | ✅ in shared core |
| `SparseDownsample(2)` | ✅ in shared core |
| `VarLenTensor` / `SparseTensor` split + `get/register_spatial_cache` | ✅ added to shared core |
| dense `nn.Conv3d(.., 2, stride=2)` in `sparse_structure_vae` | ✅ maps to `mlx.nn.Conv3d` |
| `SparseUpsample(2)` | ❌ **to do** — cache-paired inverse of a downsample |
| `SparseSpatial2Channel(2)` | ❌ **to do** — sparse pixel-shuffle, spatial→channel |
So the gap is **two ops**, both of which work through the spatial cache (they pair with a
matching downsample rather than recomputing structure). Everything else is already
covered by work done for LATO.2.
## Model surface
```
pixal3d/models/
sparse_structure_vae.py dense Conv3d — voxel structure
sparse_structure_flow.py structure flow (SS)
structured_latent_flow.py SLAT flow
sc_vaes/sparse_unet_vae.py the only file using sparse conv
```
Weights: 24.04 GB across 19 files (1.3B DiTs at 512/1024 + shape/tex decoders).
## Status
- [x] Scoped against the shared core — gap is two ops
- [x] Weights downloaded
- [ ] `SparseUpsample`, `SparseSpatial2Channel`
- [ ] Weight converter
- [ ] Model graphs
- [ ] End-to-end