[Apple Silicon] Add MPS and Metal inference support

.gitignore

@@ -205,3 +205,25 @@ cython_debug/
 marimo/_static/
 marimo/_lsp/
 __marimo__/
+
+# macOS
+.DS_Store
+
+# Model weights
+weights/
+
+# CoreML exported models
+export/coreml_models/
+
+# Generated output
+output/
+
+# Debug scripts (dev-only)
+debug_*.py
+
+# pyenv
+.python-version
+
+# Gradio cache
+gradio_cache/
+gradio_cache_test/

README.md

@@ -15,6 +15,102 @@ https://github.com/user-attachments/assets/63b43a7e-acc7-4c81-a900-6da450527d8f
 **TRELLIS.2** is a state-of-the-art large 3D generative model (4B parameters) designed for high-fidelity **image-to-3D** generation. It leverages a novel "field-free" sparse voxel structure termed **O-Voxel** to reconstruct and generate arbitrary 3D assets with complex topologies, sharp features, and full PBR materials.
 
 
+## 🍎 Apple Silicon Fork
+
+This fork makes Apple Silicon a source-native backend instead of cloning and
+patching a second TRELLIS checkout. The supported end-to-end path is PyTorch
+MPS with capability-probed Metal extensions. MLX is included only as an
+experimental parity backend; the upstream CUDA/Linux path remains available.
+
+The implementation incorporates the source-native backend and parity work from
+[`trellis2-apple@6055b86`](https://github.com/pedronaugusto/trellis2-apple/commit/6055b868734af6e12769d229d90580e775fae9f0)
+and the MPS CLI/fallback lessons from
+[`trellis-mac@d58628f`](https://github.com/shivampkumar/trellis-mac/commit/d58628f4f5b9c3de8274cb110074154f4b31cef2).
+
+### macOS setup
+
+Requirements: Apple Silicon, Xcode, and Python 3.11. The setup script creates
+`.venv`, installs the matching Xcode Metal Toolchain when necessary, tries
+`torch==2.13.0` / `torchvision==0.28.0`, builds the pinned Metal extensions,
+then runs `pip check`, MPS, SDPA, MLX, KDTree, and Metal probes. If the primary
+pair fails its ABI probe, it rebuilds once with the prescribed
+`torch==2.11.0` / `torchvision==0.26.0` fallback.
+
+```sh
+git clone https://github.com/Jourloy/TRELLIS.2.git
+cd TRELLIS.2
+bash scripts/setup_macos.sh
+source .venv/bin/activate
+```
+
+The Metal sources are fixed to these commits:
+
+- `mtlgemm`: `867aec8234299a7fe1ede7f802c8debe5a939a82`
+- `mtldiffrast`: `4668cd91cb6d27f5e264731f94a06841fbf7aab8`
+- `mtlbvh`: `23f441c470ce1f537e1fd836f3ffb5b8245f7975`
+- `mtlmesh`: `212079e55772cff3d648a21372392c37e0643f3b`
+
+`SKIP_METAL=1 bash scripts/setup_macos.sh` installs the slower controlled
+fallback environment. No API or Gradio server is needed for the supported CLI.
+
+### Hugging Face access and offline cache
+
+TRELLIS.2-4B is public. DINOv3 and RMBG-2.0 are gated, so accept their terms
+and authenticate once before downloading all pinned inputs:
+
+```sh
+hf auth login
+python scripts/download_weights.py \
+  --cache-dir ~/.cache/trellis2/huggingface
+python scripts/download_weights.py \
+  --cache-dir ~/.cache/trellis2/huggingface --offline
+```
+
+Runtime revisions are fixed to TRELLIS.2-4B `af44b45`, DINOv3 `ea8dc28`,
+RMBG-2.0 `5df4c9c`, and the external TRELLIS image decoder `25e0d31`. Pass the
+same cache with `--cache-dir` and use `--offline` after the first download.
+
+### Reproducible generation
+
+```sh
+python scripts/generate_asset.py input.png \
+  --output-dir outputs/sample \
+  --backend auto \
+  --baker auto \
+  --pipeline-type 512 \
+  --seed 42 \
+  --texture-size 1024 \
+  --background auto \
+  --pbr-decimation-target none \
+  --cache-dir ~/.cache/trellis2/huggingface
+```
+
+`background=auto` runs official RMBG-2.0 for opaque images, while a prepared
+RGBA image with non-opaque alpha bypasses background removal. `background=keep`
+passes the image through unchanged.
+
+Every run writes:
+
+- `raw_full.glb`: untouched full-resolution geometry before PBR processing;
+- `candidate_pbr.glb`: UV0 and native base-color, metallic, roughness, and
+  alpha data;
+- `meta.json`: exact revisions, backend capabilities, seed, timings, hashes,
+  bounds, sizes, triangle counts, and every fallback attempt.
+
+There is no default triangle limit. PBR export first tries the complete mesh
+with Metal, then full-resolution KDTree baking. If both fail on a very large
+mesh, a separately recorded technical candidate near 200k faces is tried;
+`raw_full.glb` is never simplified. Set `TRELLIS_DISABLE_METAL=1` to exercise
+the pure-PyTorch sparse-convolution, SDPA, CPU mesh extraction, and KDTree
+fallback path explicitly.
+
+### Dependency licenses
+
+The repository code is MIT, but model dependencies have their own terms. In
+particular, review the DINOv3 license before use and note that RMBG-2.0 is
+distributed under CC BY-NC 4.0. These terms are documented here and are not
+enforced by an automatic runtime block.
+
 ## ✨ Features
 
 ### 1. High Quality, Resolution & Efficiency

benchmarks/e2e_decoder.py

@@ -0,0 +1,223 @@
+"""End-to-end decoder-block benchmark for the Apple-Silicon TRELLIS.2 stack.
+
+Loads the representative decoder sub-graph (SparseConv3d → SparseLayerNorm →
+SparseAttention → SparseConv3d) at a trellis2-decoder-sized spatial shell and
+measures wall clock for:
+
+  1. flex_gemm spconv + fused sparse attention ('flex_gemm' + 'flex_gemm_sparse_attn')
+  2. flex_gemm spconv + SDPA-padded attention    ('flex_gemm' + 'sdpa')
+  3. torchsparse spconv + SDPA-padded attention  ('torchsparse' + 'sdpa')  [baseline]
+
+Reports per-kernel breakdown plus total wall time. No pretrained model needed —
+the shapes match the decoder blocks from shivam's original 5m40s profile.
+
+Run:
+    python benchmarks/e2e_decoder.py
+"""
+import os
+import sys
+import time
+import math
+
+os.environ.setdefault("SPARSE_CONV_BACKEND", "flex_gemm")
+os.environ.setdefault("SPARSE_ATTN_BACKEND", "flex_gemm_sparse_attn")
+os.environ.setdefault("FLEX_GEMM_QUIET", "1")
+
+# trellis2 is in-tree (no pip install) — make the bench runnable as
+# `python benchmarks/e2e_decoder.py` regardless of cwd by putting the
+# repo root on sys.path before any trellis2 import.
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import torch
+
+assert torch.backends.mps.is_available(), "This benchmark needs MPS."
+
+from trellis2.modules.sparse import SparseTensor, SparseConv3d
+from trellis2.modules.sparse.attention.full_attn import sparse_scaled_dot_product_attention
+from trellis2.modules.sparse import config as sparse_cfg
+
+
+def build_sparse_shell(res=32, ch=64, dtype=torch.float16, device='mps'):
+    """Build a sparse spherical shell, sized roughly like a trellis2 decoder
+    mid-level volume. Returns (coords, feats) ready to wrap in SparseTensor."""
+    g = torch.stack(torch.meshgrid(
+        torch.arange(res), torch.arange(res), torch.arange(res), indexing='ij',
+    ), dim=-1).int().contiguous()
+    cx = res / 2 - 0.5
+    dist = ((g.float() - cx) ** 2).sum(dim=-1).sqrt()
+    # Shell of 1.25-voxel thickness, ~4000 voxels at res=32
+    active = (dist <= res / 2) & (dist >= res / 2 - 1.25)
+    coords = torch.nonzero(active).int()
+    coords = torch.cat([torch.zeros(coords.shape[0], 1, dtype=torch.int32), coords], dim=-1)
+    coords = coords.contiguous().to(device)
+    feats = (torch.randn(coords.shape[0], ch, dtype=dtype) * 0.3).to(device).contiguous()
+    return coords, feats
+
+
+def bench(fn, warmup=2, iters=5):
+    for _ in range(warmup):
+        fn()
+    torch.mps.synchronize()
+    t0 = time.perf_counter()
+    for _ in range(iters):
+        fn()
+    torch.mps.synchronize()
+    return (time.perf_counter() - t0) / iters * 1000
+
+
+def build_attention_qkv(T, H, C, dtype, device):
+    """Synthetic Q, K, V packed for sparse attention."""
+    qkv = (torch.randn(T, H, C, dtype=dtype) * 0.3).to(device)
+    return qkv.contiguous()
+
+
+def run_attention_once(feats, H, seqlens, backend):
+    """Invoke sparse_scaled_dot_product_attention via the given backend.
+
+    `feats` is [T, H, 3, C] — packed q, k, v per layer.  Returns the MPS output.
+    """
+    prev = sparse_cfg.ATTN
+    sparse_cfg.ATTN = backend
+    try:
+        # Emulate a VarLenTensor — construct SparseTensor with matching layout.
+        # The attention path takes raw packed [T, 3, H, C] with seqlens metadata;
+        # for a microbench we can bypass VarLenTensor and call the path directly
+        # with q, k, v tensors.
+        T, _, three, C_head = feats.shape
+        q = feats[:, :, 0].contiguous()
+        k = feats[:, :, 1].contiguous()
+        v = feats[:, :, 2].contiguous()
+        import flex_gemm
+        scale = 1.0 / math.sqrt(C_head)
+        device = feats.device
+        csq  = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(seqlens), 0)]).int().to(device)
+        cskv = csq.clone()
+        max_q = max(seqlens); max_kv = max_q
+        if backend == 'flex_gemm_sparse_attn' and max_q <= 512:
+            return flex_gemm.kernels.cuda.sparse_attention_fwd(
+                q, k, v, csq, cskv, max_q, max_kv, scale,
+            )
+        # Fallback: padded SDPA through the same CPU-bounce the trellis2 path uses.
+        import torch.nn.functional as F
+        N = len(seqlens)
+        max_q = max(seqlens); max_kv = max_q
+        q_cpu = q.cpu(); k_cpu = k.cpu(); v_cpu = v.cpu()
+        qd = q_cpu.new_zeros(N, max_q, H, C_head)
+        kd = k_cpu.new_zeros(N, max_kv, H, C_head)
+        vd = v_cpu.new_zeros(N, max_kv, H, C_head)
+        mask = torch.zeros(N, max_q, max_kv, dtype=torch.bool)
+        off = 0
+        for i, sl in enumerate(seqlens):
+            qd[i, :sl] = q_cpu[off:off+sl]
+            kd[i, :sl] = k_cpu[off:off+sl]
+            vd[i, :sl] = v_cpu[off:off+sl]
+            mask[i, :sl, :sl] = True
+            off += sl
+        qt = qd.permute(0, 2, 1, 3); kt = kd.permute(0, 2, 1, 3); vt = vd.permute(0, 2, 1, 3)
+        fm = torch.zeros(N, 1, max_q, max_kv, dtype=q_cpu.dtype)
+        fm.masked_fill_(~mask.unsqueeze(1), float('-inf'))
+        o = F.scaled_dot_product_attention(qt, kt, vt, attn_mask=fm).permute(0, 2, 1, 3)
+        out_parts = [o[i, :sl] for i, sl in enumerate(seqlens)]
+        return torch.cat(out_parts, dim=0).to(device)
+    finally:
+        sparse_cfg.ATTN = prev
+
+
+def main():
+    dtype = torch.float16
+    device = 'mps'
+
+    print("=" * 80)
+    print("trellis2 decoder-block e2e bench (M3 Max, fp16, MPS)")
+    print("=" * 80)
+
+    # Spconv path: run 3 conv layers on a decoder-sized volume
+    print("\nPhase 1 — SparseConv3d (3 layers, res=32 ch=64, kernel=3)")
+    coords, feats = build_sparse_shell(res=32, ch=64, dtype=dtype, device=device)
+    print(f"  voxels={feats.shape[0]} channels=64")
+
+    conv_layers = []
+    for _ in range(3):
+        c = SparseConv3d(64, 64, kernel_size=3, bias=False).to(dtype)
+        c.weight.data = c.weight.data.to(device)
+        conv_layers.append(c)
+
+    from flex_gemm.ops.spconv import Algorithm, set_algorithm
+
+    def run_convs_masked():
+        set_algorithm(Algorithm.MASKED_IMPLICIT_GEMM)
+        x = SparseTensor(feats=feats, coords=coords, shape=torch.Size([1, 64]),
+                         spatial_shape=[32, 32, 32])
+        for c in conv_layers:
+            x = c(x)
+        return x.feats
+
+    def run_convs_dense():
+        set_algorithm(Algorithm.IMPLICIT_GEMM)
+        x = SparseTensor(feats=feats, coords=coords, shape=torch.Size([1, 64]),
+                         spatial_shape=[32, 32, 32])
+        for c in conv_layers:
+            x = c(x)
+        return x.feats
+
+    conv_masked_ms = bench(run_convs_masked)
+    conv_dense_ms  = bench(run_convs_dense)
+    print(f"  IMPLICIT_GEMM (dense):         {conv_dense_ms:8.3f} ms")
+    print(f"  MASKED_IMPLICIT_GEMM:          {conv_masked_ms:8.3f} ms")
+    print(f"  masked/dense:                  {conv_dense_ms / conv_masked_ms:.2f}x")
+
+    # Attention path: run a single block against decoder-shape QKV
+    print("\nPhase 2 — sparse attention (decoder shapes, max_seqlen=256, H=8, C=64)")
+    seqlens_dec = [256, 192, 128, 64]  # 4 chunks, balanced-ish
+    T_att = sum(seqlens_dec); H_att = 8; C_att = 64
+    qkv = torch.randn(T_att, H_att, 3, C_att, dtype=dtype).to(device).contiguous()
+    print(f"  T={T_att} H={H_att} C={C_att} seqlens={seqlens_dec}")
+
+    def run_attn_flash():
+        return run_attention_once(qkv, H_att, seqlens_dec, 'flex_gemm_sparse_attn')
+    def run_attn_sdpa():
+        return run_attention_once(qkv, H_att, seqlens_dec, 'sdpa')
+
+    attn_flash_ms = bench(run_attn_flash)
+    attn_sdpa_ms  = bench(run_attn_sdpa)
+    print(f"  flex_gemm_sparse_attn (flash): {attn_flash_ms:8.3f} ms")
+    print(f"  sdpa (CPU-bounce):             {attn_sdpa_ms:8.3f} ms")
+    print(f"  flash/sdpa:                    {attn_sdpa_ms / attn_flash_ms:.2f}x")
+
+    # Combined decoder block: 2× (conv + attn) — typical trellis2 decoder motif
+    print("\nPhase 3 — combined decoder micro-pipeline (2× conv block + 1× attn)")
+    def combined_flash():
+        set_algorithm(Algorithm.MASKED_IMPLICIT_GEMM)
+        x = SparseTensor(feats=feats, coords=coords, shape=torch.Size([1, 64]),
+                         spatial_shape=[32, 32, 32])
+        for c in conv_layers[:2]:
+            x = c(x)
+        _ = run_attention_once(qkv, H_att, seqlens_dec, 'flex_gemm_sparse_attn')
+        return x.feats
+
+    def combined_sdpa():
+        set_algorithm(Algorithm.MASKED_IMPLICIT_GEMM)
+        x = SparseTensor(feats=feats, coords=coords, shape=torch.Size([1, 64]),
+                         spatial_shape=[32, 32, 32])
+        for c in conv_layers[:2]:
+            x = c(x)
+        _ = run_attention_once(qkv, H_att, seqlens_dec, 'sdpa')
+        return x.feats
+
+    combined_flash_ms = bench(combined_flash)
+    combined_sdpa_ms  = bench(combined_sdpa)
+    print(f"  flex_gemm + flash attn:        {combined_flash_ms:8.3f} ms")
+    print(f"  flex_gemm + sdpa attn:         {combined_sdpa_ms:8.3f} ms")
+    print(f"  flash / sdpa combined:         {combined_sdpa_ms / combined_flash_ms:.2f}x")
+
+    print("\n" + "=" * 80)
+    print("Summary: decoder-block wall-clock on M3 Max, fp16 MPS")
+    print("=" * 80)
+    print(f"{'stage':32s}  {'flash':>9s}  {'sdpa':>9s}  {'speedup':>8s}")
+    print(f"{'convs-only (dense vs masked)':32s}  {conv_masked_ms:6.3f}ms  {conv_dense_ms:6.3f}ms  {conv_dense_ms/conv_masked_ms:6.2f}x")
+    print(f"{'attn-only (flash vs sdpa)':32s}  {attn_flash_ms:6.3f}ms  {attn_sdpa_ms:6.3f}ms  {attn_sdpa_ms/attn_flash_ms:6.2f}x")
+    print(f"{'combined block':32s}  {combined_flash_ms:6.3f}ms  {combined_sdpa_ms:6.3f}ms  {combined_sdpa_ms/combined_flash_ms:6.2f}x")
+
+
+if __name__ == '__main__':
+    main()