144 lines
9.4 KiB
Markdown
144 lines
9.4 KiB
Markdown
# Local 3D → full MLX + preprocessing frontend — master build plan
|
||
|
||
**For:** a fresh long-running Fable session on m3ultra (no execution-time cap).
|
||
**Author:** Opus planning pass, 2026-07-19. **Read first:** `docs/TRELLIS2_MLX_RECON.md`,
|
||
`BENCHMARKS.md` (fleet MLX matrix + rung 1/2), `HANDOFF_HY3D_MLX.md`, `CORRIDORKEY.md`.
|
||
**Prime directive:** parity-first. Every kernel/port change ships with a PT-parity test
|
||
(≤1e-4 rel) + a BENCHMARKS.md row (per-stage time + `mx.get_peak_memory()`). No silent
|
||
regressions — that discipline is why hunyuan3d-mlx and corridorkey-mlx worked.
|
||
|
||
**Two goals, independent, run in parallel:**
|
||
- **G1 — the models go fully MLX + fast** (TRELLIS.2 port from scratch-ish; Hunyuan tune).
|
||
- **G2 — a local preprocessing frontend** that turns any messy photo into the input the
|
||
models want. Model-agnostic; helps every generation from here on.
|
||
|
||
---
|
||
|
||
## Ground truth (measured, don't re-derive)
|
||
|
||
- Local TRELLIS.2 = `vendor/trellis-mac` (shivampkumar fork), **torch-MPS, NOT MLX**.
|
||
124.8s/gen on m3ultra (1024-cascade, tex2048), 20.8GB peak. Weights `TRELLIS.2-4B` 14GB.
|
||
- Stage breakdown of the 124.8s: **~46s diffusion sampling** (padded SDPA attention),
|
||
**~50s decode + mesh extract** (incl. `aten::segment_reduce` → **CPU fallback**),
|
||
**19s Metal PBR bake**. Decoder emits **2.79M faces** pre-simplify.
|
||
- Sparse-conv backend already on the Metal fast path (`flex_gemm`); worth **1.72×** vs the
|
||
pure-torch fallback (rung 1). Metallib loads on macOS 26.5.
|
||
- Hunyuan3D-2.1 = `vendor/hunyuan3d-mlx` (dgrauet), **already MLX**, ~260s m3ultra, runs on M1.
|
||
- Quality gap vs fal = dark shattered patches; density is NOT the cause (rung 2: 200K/800K/2.79M
|
||
all show it) → **texture bake or decoder attrs**. KDTree-baker A/B pending (`TRELLIS2_FORCE_KDTREE`).
|
||
- Fleet MLX matrix (BENCHMARKS.md): bf16 fine except M1 (~20% slower → ship fp16 on M1);
|
||
**pad attention head_dim→64** (up to 4.7× M5, 1.5× M3U, ~0 on M1); `mx.compile` on fixed shapes.
|
||
- Knobs already added to `vendor/trellis-mac/generate.py`: `TRELLIS2_MAX_BAKE_FACES` (0=uncapped),
|
||
`TRELLIS2_FORCE_KDTREE` (A/B the baker). Bench artifacts in `~/Documents/trellis2-bench/`.
|
||
|
||
---
|
||
|
||
## G1 · WORKSTREAM A — TRELLIS.2 → MLX
|
||
|
||
### A0. Scout the existing MLX efforts (do FIRST — decides adopt vs build)
|
||
- Vendor `pedronaugusto/trellis2-apple` (has an `mlx_backend/`) and upstream **PR #175**
|
||
(MPS+Metal + experimental `mlx` flag, 28 tests green on M4 Max). Fork both to Gitea as
|
||
`monster/trellis2-mrp-*` (house pattern; keep upstream remote for rebases).
|
||
- Install each, run ONE gen on the anatomy bench image, seed 42, 1024-cascade. Compare to the
|
||
124.8s MPS baseline **and** to fal quality (renders already in `trellis2-bench/`).
|
||
- **Decision gate:** if trellis2-apple's MLX path is correct + faster → we ADOPT and jump to A3
|
||
(patch, not rebuild). If it's partial/slower/buggy → we BUILD (A1→A2), using it + the
|
||
corridorkey `prompts/` 6-phase template as reference.
|
||
|
||
### A1. Scaffold the MLX port (if building)
|
||
- New repo `monster/trellis2-mlx`. Mirror the corridorkey-mlx phased layout: `prompts/` phase docs,
|
||
`src/`, `tests/test_parity.py`, `scripts/convert_weights.py`.
|
||
- Weight conversion: TRELLIS.2-4B safetensors → MLX. Reuse the hunyuan pattern
|
||
(`hy3dpaint/utils/convert_realesrgan.py`): torch load → key remap → **transpose Conv (O,I,H,W)→(O,H,W,I)** → save. fp16 (M1-safe); publish converted weights to HF/Gitea so re-runs skip conversion.
|
||
- Build a PT reference harness first (dump per-stage activations from the working MPS build as
|
||
fixtures) — every MLX module validates against these.
|
||
|
||
### A2. The kernel rewrites (the real work — this is where the days go)
|
||
Priority = by measured cost. Each = an `mx.fast.metal_kernel` + parity test + bench row.
|
||
1. **Sparse-conv (FlexGEMM → fused Metal).** THE crux and the ~10× lever. Triton has no Metal
|
||
target; implement gather→GEMM→scatter as a fused Metal kernel. Reference: the existing
|
||
`mtlgemm` Metal port in trellis-mac/deps and `backends/conv_none.py` (the slow pure-torch
|
||
version) for correct semantics. MLX has no sparse tensor type — hand-roll SparseTensor as
|
||
coord-list + hash (same as CUDA/MPS side).
|
||
2. **Varlen/sparse attention (padded SDPA → fused).** ~46s of the run. Kill the pad-to-maxlen
|
||
waste; pad head_dim→64 for the fused fast path (fleet matrix). This + #1 close most of the gap.
|
||
3. **`segment_reduce` MPS→CPU fallback → MLX-native.** Currently silently on CPU each step;
|
||
pure MLX scatter-add removes a real tax in the decode path.
|
||
4. **o-voxel hashmap + dual-grid mesh extract.** Keep the Metal/CPU fork if it's not hot; port to
|
||
MLX only if profiling says so.
|
||
Leave the PBR bake on the existing Metal path (`mtldiffrast`) unless G1-C convicts it.
|
||
|
||
### A3. Fleet-tune + the prize
|
||
- Apply corridorkey ablation moves across the port: head_dim→64, `mx.compile` fixed shapes,
|
||
sdpa gating by GPU generation, tiled-vs-full sweeps. Nightly perfcheck like corridorkey.
|
||
- **Prize: M1 Ultra compatibility.** MLX-native runs where torch-MPS/bf16 is shaky → a *second*
|
||
free TRELLIS box (297 GB/s, 128GB). Validate a full gen on `johnking@100.91.239.7`.
|
||
- Wire `trellis_mac`→MLX into MODELBEAST `nodes.json` (m3ultra primary, m1ultra fp16 second) and
|
||
add `trellis2` to MESHGOD's `local/` model list beside `hunyuan3d_mlx`.
|
||
|
||
## G1 · WORKSTREAM B — Hunyuan3D-2.1-MLX optimization (already MLX; just tune)
|
||
- It's a *tune*, not a port. Apply the same corridorkey wins to `vendor/hunyuan3d-mlx`:
|
||
pad attention head_dim→64, `mx.compile` on fixed-shape blocks (DiT denoiser, VAE), sdpa gating.
|
||
- Check `mlx_arsenal` rasterizer (`mesh_render_mlx.py`, 1137 LOC) for the same fixed-shape
|
||
compile opportunity. Bench each on the fleet matrix; expect corridorkey-scale wins (1.1–1.5×).
|
||
- Low risk, high certainty — good "warm-up" task while A0 scouting runs.
|
||
|
||
## G1 · WORKSTREAM C — quality (rung 2 close-out; feeds BOTH MPS + MLX)
|
||
- **KDTree-baker A/B** (`TRELLIS2_FORCE_KDTREE=1`, running now). If its surface is clean → the
|
||
**Metal baker (mtldiffrast/mtlbvh texel→voxel sampling) is the dark-patch culprit** → fix or
|
||
default-swap the baker. If still dirty → decoder voxel attrs are suspect → investigate upstream.
|
||
- Keep `TRELLIS2_MAX_BAKE_FACES` (geometry win, free).
|
||
- Add a **hole-fill pass** (CuMesh replacement) pre-bake — Blender fill or a voxel-remesh.
|
||
- Loud assertion if `flex_gemm` metallib ever fails to load (silent → +90s/gen forever).
|
||
|
||
---
|
||
|
||
## G2 · WORKSTREAM D — image preprocessing frontend ("best chance of success")
|
||
|
||
North star already written in `meshgod/config.py`: *centered · plain neutral light-gray bg ·
|
||
even shadowless studio light · neutral symmetric A-pose*. The job = transform any user photo → that.
|
||
Model-agnostic; helps TRELLIS.2 AND Hunyuan equally; no fal/network.
|
||
|
||
### D1. `prep_local()` — the front line (highest ROI, pure glue)
|
||
- RMBG-2.0 (local `rmbg` venv) or BiRefNet matte → **composite on neutral gray (NOT alpha** —
|
||
config's own spec) → bbox crop → center → square-pad → upscale to ≥1024.
|
||
- Wire as a LOCAL replacement for MESHGOD's current `clean_bg` (today it's a fal call). Respect the
|
||
learned order: **upscale FIRST, matte LAST** (SeedVR strips alpha).
|
||
- ⚠️ **Licensing:** RMBG-2.0 is CC-BY-NC → local/personal only. Ship-commercial path = **BiRefNet
|
||
(Apache)**. Expose both; default BiRefNet for anything headed to a game.
|
||
|
||
### D2. Lighting/exposure normalization
|
||
- Auto-levels + gray-world WB before recon. TRELLIS bakes light into albedo → even input = cleaner
|
||
texture, and directly fights the dark-patch failure mode. Trivial PIL; measurable.
|
||
|
||
### D3. Multi-view synthesis — the RIGHT fix for touching limbs
|
||
- Single-view can't separate an arm fused to a torso in silhouette. Generate back + ¾ views (image
|
||
model) → feed **multi-view reconstruction** (both models support it; MESHGOD already wires the MV
|
||
path). The extra angle disambiguates. Prototype: photo → gen back view → MV recon → compare.
|
||
|
||
### D4. Input QA gate (save wasted gens)
|
||
- Pre-flight heuristic/VLM (local Ollama VLM or DINOv3-embedding rules) flags too-small, busy-bg,
|
||
cropped-limb, multi-subject BEFORE spending a generation. Cheap insurance.
|
||
|
||
### D5. T-pose normalization = POST-process via MIRPAMO (NOT a preprocess)
|
||
- Reposing a 2D human hallucinates. Correct pipeline: reconstruct in-pose → **MIRPAMO auto-rig +
|
||
retarget to T/A-pose**. Design-doc it; tie in MOCAPGOD/HSMR (SMPL fit, on m1) as the body prior.
|
||
|
||
---
|
||
|
||
## Suggested execution order for the long session
|
||
1. **G1-C** close-out (read KDTree result, act) — unblocks the quality question. *(minutes)*
|
||
2. **G1-B** Hunyuan tune — certain win, warms up the parity/bench rhythm. *(hours)*
|
||
3. **G2-D1+D2** preprocessing front line — compounding ROI on every future gen. *(hours)*
|
||
4. **G1-A0** scout MLX ports — the adopt-vs-build decision gate. *(hour)*
|
||
5. **G1-A1→A3** the TRELLIS.2 MLX port + kernels — the multi-day main event. *(days)*
|
||
6. **G2-D3/D4/D5** as reach goals.
|
||
|
||
## Definition of done
|
||
- TRELLIS.2 runs MLX on BOTH Ultras, faster than today's 124.8s MPS, parity ≤1e-4, in `nodes.json`.
|
||
- Hunyuan tuned with a measured fleet-matrix win.
|
||
- Dark-patch quality gap root-caused and closed (or a clear upstream bug filed).
|
||
- `prep_local()` live as MESHGOD's default local pre-pass; BiRefNet for commercial, multi-view
|
||
path proven on a touching-limbs case.
|
||
- Every step has a BENCHMARKS.md row. Upstreamable kernel work offered to PR #175.
|