modelbeast/docs/TRELLIS2_HY3D_MLX_BUILD_PLAN.md

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.