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:
commit
80991a8fe9
6
.gitignore
vendored
Normal file
6
.gitignore
vendored
Normal file
@ -0,0 +1,6 @@
|
||||
.venv
|
||||
upstream/
|
||||
ckpt/
|
||||
weights/
|
||||
__pycache__/
|
||||
*.pyc
|
||||
29
CLAUDE.md
Normal file
29
CLAUDE.md
Normal 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
56
README.md
Normal 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
|
||||
Loading…
Reference in New Issue
Block a user