2ede6655d9
The status list still said the pipeline was unwired. It runs end to end now: proj conditioning (all four extractors at corr 1.00000000), the NAF branch without natten, the full 32^3 -> refine -> 64^3 cascade at silhouette IoU 0.969, cleanup, and the MODELBEAST operator. Textures remain the one open item. Also records the decimation table, because the floor is a real constraint and not obvious: 500k faces is effectively lossless (IoU 0.965) but ~214k cannot be beaten, and reaching it costs fidelity (0.823). Cause is the ~180k boundary edges the dual grid emits for open surfaces - quadric decimation will not collapse them at any setting, and o_voxel's remesh ran >20min on 214k faces before being killed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
165 lines
8.1 KiB
Markdown
165 lines
8.1 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.
|
|
- [x] DiT blocks: RoPE (stored complex phases), qk_rms_norm, AdaLN modulation, proj conditioning
|
|
- [x] **All four flow models verified against upstream at correlation 1.00000000**
|
|
— `ss_flow` (max diff 1.2e-5) and `slat_flow` (9.3e-6), 700/700 params each,
|
|
on the real 1.3B checkpoints
|
|
- [x] `SparseConvNeXtBlock3d`, `SparseResBlockC2S3d`, `spatial2channel`/`channel2spatial`
|
|
— all in the shared core, round-trip and selective-growth tested
|
|
- [x] **`ss_dec` verified at correlation 1.00000000** (74/74), completing the whole
|
|
structure stage: image -> ss_flow -> latent -> ss_dec -> 64^3 occupancy grid
|
|
- [x] `shape_dec` / `tex_dec` — both load complete (292/292, 284/284) and run.
|
|
**Behaviourally** checked only; see the verification note below
|
|
- [x] **Proj conditioning** — camera back-projection, an exact `grid_sample`
|
|
equivalent, and DINOv3 (torch/MPS). All four extractor stages diffed against
|
|
upstream at correlation 1.00000000
|
|
- [x] **NAF high-res branch without natten** — natten is unusable on Apple Silicon
|
|
(no `libnatten`; `flex-fna` is CPU-only, rejects asymmetric head dims, and was
|
|
OOM-killed at 512). Replaced by an exact reduction to a clamped 9x9 low-res
|
|
neighbourhood, verified against natten at 7.15e-07
|
|
- [x] **The full cascade** — 32^3 -> LR SLAT -> coord refinement -> 64^3 HR SLAT ->
|
|
Flexible Dual Grid at 1024^3 -> GLB, at **silhouette IoU 0.969**
|
|
- [x] **Mesh cleanup** — weld, strip floaters, iterative decimation (500k faces at
|
|
IoU 0.965; ~214k is a hard floor, see below)
|
|
- [x] **MODELBEAST operator** `pixal3d_mlx` (geometry only)
|
|
- [ ] Texture stage (tex SLAT + o_voxel PBR bake) — the GLB is currently untextured
|
|
|
|
## Running it
|
|
|
|
```bash
|
|
python scripts/image_to_mesh.py IMAGE -o out.glb --target-faces 500000
|
|
```
|
|
|
|
Exits non-zero below `--min-iou` (default 0.85): a run that completes with a
|
|
reconstruction that does not track the input has failed, even though nothing raised.
|
|
|
|
### Mesh quality, measured
|
|
|
|
| face budget | result | silhouette IoU |
|
|
|---|---|---|
|
|
| raw decoder output | 7,996,876 | 0.969 |
|
|
| 500,000 | 499,984 | **0.965** |
|
|
| 200,000 / 100,000 / 20,000 | 214,322 (floor) | 0.823 |
|
|
|
|
**~214k is a hard floor.** The Flexible Dual Grid emits ~180,000 boundary edges for
|
|
open surfaces and quadric decimation will not collapse those — no `target_reduction`
|
|
or `agg` setting changes it, and a single `fast_simplification` call additionally
|
|
refuses to reduce past ~4.4% of its input (hence the iterative loop). Going lower
|
|
needs a remesh; o_voxel's ran >20 minutes on 214k faces before being killed, so it is
|
|
not currently a practical route. **500k is effectively lossless — use that.**
|
|
|
|
|
|
## Model status
|
|
|
|
| model | params | verification |
|
|
|---|---|---|
|
|
| `ss_flow` | 700/700 | **corr 1.00000000** vs upstream |
|
|
| `slat_flow` x3 | 700/700 | **corr 1.00000000** vs upstream |
|
|
| `ss_dec` | 74/74 | **corr 1.00000000** vs upstream |
|
|
| `shape_dec` | 292/292 | behavioural only — no oracle possible |
|
|
| `tex_dec` | 284/284 | behavioural only — no oracle possible |
|
|
|
|
The split is not arbitrary: the first five contain no sparse convolution, so upstream
|
|
runs on CPU torch and can be diffed directly. The two decoders do use it, spconv has no
|
|
Metal build, and so there is nothing to diff against. Their blocks are individually
|
|
tested, and the assembled graphs are checked for complete weight mapping, selective
|
|
growth, correct scale and a vertex head inside its valid band — but that is weaker
|
|
evidence than a correlation and should be read that way.
|
|
|
|
## Numerical verification
|
|
|
|
The flow models contain no sparse convolution, which means **upstream runs on CPU torch
|
|
here** — swap flash-attn for `F.scaled_dot_product_attention` and it loads the real
|
|
checkpoint and runs. So unlike the sparse path (where spconv is uninstallable and the
|
|
oracle had to be hand-written), these are diffed against upstream directly:
|
|
|
|
```
|
|
MLX : mean +0.18490 std 0.86841
|
|
UPSTREAM : mean +0.18490 std 0.86841
|
|
max abs diff 1.216e-05 correlation 1.00000000
|
|
```
|
|
|
|
Getting there required finding three bugs that **weight-key matching could not catch** —
|
|
the loader reported a perfect 700/700 with 0 missing and 0 unmapped through all of them:
|
|
|
|
1. **A parameterless final `LayerNorm`** between the last block and `out_layer`. It has
|
|
no weights, so it leaves no trace in the checkpoint. Without it the output was ~200x
|
|
too large (std 187 vs 0.87).
|
|
2. **`rope_phases` is complex64**, built with `torch.polar`. The rotation is a complex
|
|
multiply and cos/sin are the phase's real/imaginary parts — taking `cos()` of a
|
|
complex phase is meaningless.
|
|
3. **qk RMS norm is applied BEFORE RoPE, not after.** They do not commute. Reversed, a
|
|
single block still correlated 0.9998 with upstream; over 30 blocks that compounds to
|
|
0.84. This one is invisible without an oracle.
|