pixal3d_mrp_mlx/README.md
m3ultra bf8b1a35d5 Close all four open items: MoGe camera, manifold remesh, winding, UV bake
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>
2026-08-03 20:39:45 +10:00

192 lines
9.5 KiB
Markdown
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

# 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 2531° 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.