[Apple Silicon] Add MPS and Metal inference support

[Jourloy]

Summary

This PR adds source-native Apple Silicon support to TRELLIS.2.

On arm64 macOS, the automatic backend selects PyTorch MPS together with capability-probed Metal extensions. The existing CUDA/Linux route remains available, while macOS receives explicit fallbacks for operations that are unavailable or unstable on MPS.

This addresses the Apple Silicon / MPS portion of #74. It intentionally does not close the AMD or Intel parts of that issue.

What changed

• Add automatic Apple Silicon backend resolution:
  • PyTorch MPS for model inference
  • flex_gemm Metal sparse convolution and sparse attention
  • mtldiffrast Metal rasterization
  • mtlmesh / mtlbvh mesh and BVH operations
• Add capability probes before enabling individual Metal components.
• Add controlled fallback paths:
  • pure-PyTorch sparse convolution
  • PyTorch SDPA attention
  • CPU mesh extraction
  • KDTree texture baking
• Use SDPA for dense attention on macOS while preserving the upstream flash_attn default on other platforms.
• Add actionable diagnostics for MPS OOM, empty meshes, and Metal watchdog/BVH failures.
• Add a reproducible Python 3.11 macOS setup:
  • primary pair: PyTorch 2.13.0 / torchvision 0.28.0
  • ABI fallback: PyTorch 2.11.0 / torchvision 0.26.0
  • pinned Metal extension revisions
  • pip check, MPS, SDPA, MLX, KDTree, raster, and BVH probes
• Pin the runtime model revisions and support strict offline loading from a dedicated Hugging Face cache.
• Preserve authored RGBA transparency; opaque inputs still use the official RMBG-2.0 model.
• Add a CLI for reproducible image-to-3D generation with:
  • auto, mps, and experimental mlx backend selection
  • Metal/KDTree baker selection
  • 512, 1024, and 1024 cascade pipelines
  • deterministic seeds and configurable texture size
  • no default triangle limit
• Produce reproducible artifacts:
  • raw_full.glb
  • candidate_pbr.glb
  • meta.json with revisions, timings, hashes, bounds, backend probes, and fallback attempts
• Preserve generated UV0, base color, metallic, roughness, and alpha during PBR export.
• Add an experimental MLX backend with small numerical parity tests. End-to-end acceptance still uses PyTorch MPS.

The PBR fallback order is:

1. full-resolution Metal bake
2. full-resolution KDTree bake
3. technical ~200k-face candidate with Metal
4. technical ~200k-face candidate with KDTree

raw_full.glb is never decimated, and every fallback is recorded in meta.json.

Validation

Tested on:

• Apple M4 Max, 36 GB unified memory
• macOS 26.4
• Python 3.11.15
• PyTorch 2.13.0
• torchvision 0.28.0

Automated checks:

• 28 passed in the pytest suite
• standalone Metal sparse-convolution integration smoke passed
• dense, masked, and production split-K Metal kernels produced equivalent output
• MPS, SDPA, Metal sparse attention, raster, BVH, KDTree, and MLX probes passed
• strict offline model-cache verification passed
• pip check and compileall passed
• simulated Metal BVH failure selected the expected full-resolution KDTree fallback

End-to-end generation:

• pipeline: 512
• seed: 42
• texture size: 1024
• input: repository example RGBA image
• raw_full.glb: 3,090,522 triangles
• candidate_pbr.glb: 3,025,626 triangles
• first full-resolution Metal PBR attempt succeeded
• technical decimation was not used
• wall time: 1004.60 seconds
• maximum RSS: 17.82 GB
• completed without OOM

The resulting glTF 2.0 files were validated for non-empty geometry, UV0, vertex normals, consistent bounds, and embedded base-color, metallic-roughness, and alpha textures.

The official RMBG-2.0 path was also tested separately with an opaque input and produced a non-trivial alpha mask.

Compatibility and known limitations

• End-to-end testing was performed on Apple M4 Max only.
• The CUDA/Linux route is retained but was not end-to-end tested on CUDA hardware in this environment.
• PyTorch currently runs segment_reduce through its CPU fallback on MPS. This is functional but affects performance.
• MLX support is experimental and currently covered by small parity tests rather than end-to-end acceptance.
• DINOv3 and RMBG-2.0 require accepted Hugging Face access terms and authentication.
• The macOS setup currently depends on pinned Metal extension forks.
• This is a large change. I am happy to split the core MPS/Metal backend, setup/CLI, and experimental MLX work into separate PRs if that is easier to review.

Credits

This work builds on and preserves the authorship of the Apple backend work from:

• pedronaugusto/trellis2-apple@6055b86

The CLI diagnostics and fallback design were also informed by:

• shivampkumar/trellis-mac@d58628f

[Jourloy]

@microsoft-github-policy-service agree

[xocialize]

Nice to see the Apple Silicon path getting upstreamed — thanks for carrying the pedronaugusto/trellis2-apple lineage forward with credit.

Heads up on a small overlap: I opened two tiny fixes against main (#173, #174) that branch from the same 6055b86 you credit, and this PR touches both spots without carrying the fix, so they'll conflict at merge and — more importantly — the underlying bugs would survive:

• from_pretrained (Fix KeyError in pipelines.from_pretrained when it is the first touch of the module #174): you edit return globals()[config['name']].from_pretrained(path) to add **hub_kwargs, but keep the globals()[...] lookup. That bypasses the module's lazy __getattr__, so it raises KeyError: 'Trellis2ImageTo3DPipeline' whenever from_pretrained is the first touch of the pipelines module (nothing has populated globals() yet). One-char-ish fix: getattr(sys.modules[__name__], config['name']).

• texturing keep-existing-UVs (Texturing: make the keep-existing-UVs path reachable (and non-crashing once it is) #173): you rewrite postprocess_mesh (.cuda() → .to(self.device)) but leave normals = mesh.vertex_normals and don't carry visual=mesh.visual through preprocess_mesh. Without the visual= pass-through, run() normalizes the mesh before postprocess_mesh sees it and the keep-existing-UVs branch is unreachable — and the moment it is reachable, the in-place glTF axis swap on the read-only vertex_normals cache crashes (hence the np.array() copies). We verified the fixed path re-textures with bit-identical UVs (max delta 0.0) in production.

Both are 1–8 line changes. Happy to have you fold them into this PR (I don't need the attribution), or I can rebase #173/#174 on top of yours once one of them lands — whatever's least friction for you. No urgency on our end.

[Jourloy]

@xocialize Thanks for the detailed heads-up. I’ve folded both fixes into #175: the lazy from_pretrained resolution fix, plus the visual/UV pass-through and copied mesh arrays before the in-place transforms.