THE DECIMATION FLOOR WAS MISDIAGNOSED. I attributed it to ~180k boundary edges. It is non-manifold edges. Measured on the shipped 500k mesh: boundary edges 32,370 NON-MANIFOLD 81,112 <- the actual blocker, 2.5x more Quadric decimation cannot collapse an edge shared by more than two faces. Upstream's own fix (fill_holes, via CUDA-only cumesh) targets boundaries and caps at max_hole_perimeter=3e-2, so it was never going to help: trimesh's equivalent moved boundaries 32,370 -> 30,990 and the floor only 214k -> 210k. That falsified it. --manifold: voxelise -> fill -> marching cubes. Removes BOTH classes at once and so closes three of the four items in one change: as shipped 499,984 faces bnd 32,370 nonmani 81,112 watertight=F IoU 0.969 remeshed 1,178,142 faces bnd 0 nonmani 0 watertight=T IoU 0.949 -> 20k 19,998 faces bnd 0 winding consistent IoU 0.956 25x smaller, fully manifold, consistent winding, for 1.3% silhouette IoU. Lossy by design - it gives up the dual grid's open-surface representation - so it is opt-in. UV BAKE is unblocked by the same change: its cost is driven by face count, not by remesh. 5.0s at 20k faces against >20min at 214k. No longer offline-only when paired with manifold. THE SCALING TRAP, worth knowing: marching_cubes returns vertices in VOXEL INDEX space. Translating without apply_scale(pitch) leaves the mesh ~292x too large. It still exports and renders as a plausible object; it silhouettes at IoU 0.08. That is how it was caught. MoGe-2 CAMERA is now wired and is the default, matching upstream; --fixed-fov keeps the old constant. It runs once per image in torch/MPS, ~0.4s after load. Reporting this one straight: it did NOT improve the samples. On 1_img, fixed 49.1 deg scored 0.893 and MoGe's 29.7 deg scored 0.883. Two caveats keep it as the default anyway - the silhouette metric projects with the SAME FOV used to generate, so a wrong-but-consistent camera can still score well and the metric cannot fully arbitrate camera correctness; and the bundled samples are synthetic renders, not the photographs MoGe reads. Real photos are the intended input here, and upstream estimates too. But the constant is one flag away and the measurement is on record rather than assumed. Operator gains manifold, divisions, fixed_fov. README and PROFILE.md corrected where they repeated the boundary-edge claim. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
192 lines
9.5 KiB
Markdown
192 lines
9.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.
|
||
- [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 |
|
||
| **`--manifold` → 20,000** | **19,998** | **0.956** |
|
||
|
||
**Without `--manifold` there is a hard floor around 214k.** The cause is **non-manifold
|
||
edges (81,112)**, NOT boundary edges (32,370) — an earlier note here said the opposite.
|
||
Quadric decimation cannot collapse an edge shared by more than two faces, and no
|
||
`target_reduction` or `agg` setting changes that. Filling holes barely moved either
|
||
number (boundaries 32,370 → 30,990; floor 214k → 210k), which is what ruled the
|
||
boundary theory out.
|
||
|
||
Those non-manifold edges are the Flexible Dual Grid working as designed — it
|
||
represents open and non-manifold surfaces deliberately.
|
||
|
||
**`--manifold` voxel-remeshes past it**: watertight, zero non-manifold edges,
|
||
consistent winding, decimates to **20k faces for ~1.3% silhouette IoU** (0.969 →
|
||
0.956). Lossy — it gives up the open-surface representation and softens sharp
|
||
features — so it is opt-in. It also makes the **UV bake practical: 5.0s at 20k faces**
|
||
against >20 minutes at 214k.
|
||
|
||
Use `--manifold --target-faces 20000` for game-ready assets; plain `--target-faces
|
||
500000` when you want maximum fidelity and will clean up in Blender.
|
||
|
||
## Camera
|
||
|
||
FOV is estimated per image with **MoGe-2** by default (upstream's behaviour, ~0.4s);
|
||
`--fixed-fov` uses the 0.8576 rad constant. Everything downstream is placed by this,
|
||
so a wrong FOV reconstructs at the wrong depth scale and fails silently.
|
||
|
||
Measured honestly: on the bundled sample renders MoGe estimates 25–31° against the
|
||
49.1° constant, and scored **marginally WORSE** (0.883 vs 0.893 on `1_img`). Two
|
||
caveats on that comparison — the silhouette metric projects with the same FOV used to
|
||
generate, so a wrong-but-consistent camera can still score well; and the samples are
|
||
synthetic renders, not the photographs MoGe was trained to read. Estimation stays the
|
||
default for upstream parity and because real photos are the intended input, but the
|
||
constant is one flag away.
|
||
|
||
|
||
## 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.
|