e23233731b
Converter is light because Pixal3D ships safetensors with sibling .json configs, so there is no architecture to infer and no pickle to unpack. Only rank-5 tensors are rearranged; dtype is preserved (upcasting fp16->fp32 doubled 24GB for no benefit) and the KRSC remap is verified a pure permutation of values. Measured: the four flow models (~20GB) have ZERO 5-D tensors - pure transformers that never touch sparse conv. shape_dec/tex_dec carry 40 KRSC kernels each, ss_dec 20 dense Conv3d, and classify_5d separates them correctly on the real weights. Corrects the earlier 'gap is two ops' claim: that was right about modules/sparse but undercounted the model blocks. The configs show all four flow models need RoPE, qk_rms_norm and AdaLN modulation, and the decoders need SparseConvNeXtBlock3d and SparseResBlockC2S3d.
78 lines
3.5 KiB
Markdown
78 lines
3.5 KiB
Markdown
# 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 |
|
|
|
|
Both of those are now **done** in the shared core.
|
|
|
|
### Correction to the earlier scope
|
|
|
|
"the gap is two ops" was accurate about `modules/sparse/` — the sparse *primitives*. It
|
|
undercounted the **model blocks**, which the configs revealed:
|
|
|
|
| Still needed | Where |
|
|
|---|---|
|
|
| RoPE positional embedding | all 4 flow models (`pe_mode: "rope"`) — but `rope_phases` ships as a **stored tensor**, so phases are precomputed, not derived |
|
|
| `qk_rms_norm` on q and k | all 4 flow models |
|
|
| AdaLN modulation (`share_mod: true`) | all 4 flow models |
|
|
| `image_attn_mode: "proj"` conditioning | all 4 flow models |
|
|
| `SparseConvNeXtBlock3d` | shape_dec, tex_dec |
|
|
| `SparseResBlockC2S3d` (channel↔spatial) | shape_dec, tex_dec — uses `SparseSpatial2Channel` |
|
|
|
|
Offsetting that, a genuine simplification the tensors revealed: **the four flow models
|
|
(~20 GB, the bulk of the download) contain ZERO 5-D tensors.** `ss_flow` and the three
|
|
`slat_flow` DiTs are pure transformers — they never touch sparse convolution, so they
|
|
need none of the sparse core, just DiT blocks.
|
|
|
|
## 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
|
|
- [x] Weights downloaded (24 GB) — each ships a sibling `.json` with the exact config,
|
|
so unlike LATO.2 there is no architecture to infer
|
|
- [x] `upsample` (masked) + `downsample(mode=)` landed in the shared core
|
|
- [x] **Weight converter** — decoders remap KRSC→`[K³,in,out]`, dtype preserved,
|
|
remap verified a pure permutation. Flow models pass through untouched.
|
|
- [ ] DiT blocks: RoPE (stored phases), qk_rms_norm, AdaLN modulation
|
|
- [ ] `SparseConvNeXtBlock3d`, `SparseResBlockC2S3d`, `SparseSpatial2Channel`
|
|
- [ ] Model graphs
|
|
- [ ] End-to-end
|