add mlx backend for apple silicon with metal gpu acceleration

.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,10 @@ 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 adds an **MLX backend** for native Apple Silicon (M-series) inference, with Metal GPU acceleration for mesh postprocessing via `mtldiffrast`, `cumesh`, and `flex_gemm`. The original CUDA pipeline is fully preserved. See `mlx_backend/` for details and `requirements_macos.txt` for macOS dependencies.
+
 ## ✨ Features
 
 ### 1. High Quality, Resolution & Efficiency

api_models.py

@@ -0,0 +1,35 @@
+"""Pydantic request/response models for the Trellis2 API server."""
+from typing import Optional
+from pydantic import BaseModel, Field
+
+
+class GenerateRequest(BaseModel):
+    """Request to generate a 3D model from an image."""
+    image: str = Field(..., description="Base64-encoded image (PNG/JPEG)")
+    seed: int = Field(default=42, description="Random seed")
+    pipeline_type: str = Field(
+        default="1024_cascade",
+        description="Pipeline type: 512, 1024, 1024_cascade, 1536_cascade",
+    )
+    output_path: Optional[str] = Field(default=None, description="Save GLB to this path (in addition to returning base64)")
+    decimation_target: int = Field(default=1000000, description="Target face count for simplification")
+    texture_size: int = Field(default=2048, description="Texture resolution")
+    remesh: bool = Field(default=False, description="Whether to remesh the output")
+    steps: Optional[int] = Field(default=None, description="Number of sampler steps (default: 12, lower = faster)")
+    guidance_strength: Optional[float] = Field(default=None, description="Guidance strength for structure/shape samplers (default: 7.5)")
+    texture_guidance: Optional[float] = Field(default=None, description="Guidance strength for texture sampler (default: 1.0 = OFF)")
+
+
+class GenerateResponse(BaseModel):
+    """Response containing the generated 3D model."""
+    glb: str = Field(..., description="Base64-encoded GLB file")
+    vertices: int = Field(..., description="Number of vertices in the output mesh")
+    faces: int = Field(..., description="Number of faces in the output mesh")
+    generation_time: float = Field(..., description="Generation time in seconds")
+
+
+class HealthResponse(BaseModel):
+    """Health check response."""
+    status: str = "ok"
+    backend: str = ""
+    weights_loaded: bool = False

api_server.py

@@ -0,0 +1,157 @@
+"""
+FastAPI server for Trellis2 image-to-3D generation (MLX backend).
+
+- POST /generate — base64 image → GLB
+- GET /health — server status
+
+Usage:
+    python api_server.py --weights weights/TRELLIS.2-4B --port 8082
+"""
+import os
+import io
+import base64
+import time
+import tempfile
+import argparse
+from contextlib import asynccontextmanager
+
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+import uvicorn
+from PIL import Image
+
+from api_models import GenerateRequest, GenerateResponse, HealthResponse
+from mlx_backend import setup_logging
+
+# Global pipeline instance
+pipeline = None
+
+
+@asynccontextmanager
+async def lifespan(app):
+    global pipeline
+    from mlx_backend.pipeline import create_mlx_pipeline
+
+    weights = os.environ.get("TRELLIS2_WEIGHTS", "weights/TRELLIS.2-4B")
+    pipeline = create_mlx_pipeline(weights_path=weights)
+    yield
+    pipeline = None
+
+
+app = FastAPI(title="Trellis2 MLX API", version="1.0.0", lifespan=lifespan)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+@app.get("/health", response_model=HealthResponse)
+async def health():
+    if pipeline is None:
+        return HealthResponse(status="loading")
+    return HealthResponse(
+        status="ok",
+        backend="mlx",
+        weights_loaded=True,
+    )
+
+
+@app.post("/generate", response_model=GenerateResponse)
+async def generate(request: GenerateRequest):
+    if pipeline is None:
+        raise HTTPException(status_code=503, detail="Pipeline not ready")
+
+    t_start = time.time()
+
+    # Decode input image
+    try:
+        image_bytes = base64.b64decode(request.image)
+        image = Image.open(io.BytesIO(image_bytes))
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Invalid image: {e}")
+
+    # Build per-sampler overrides (upstream defaults: structure/shape=7.5, texture=1.0)
+    structure_shape_overrides = {}
+    if request.steps is not None:
+        structure_shape_overrides['steps'] = request.steps
+    if request.guidance_strength is not None:
+        structure_shape_overrides['guidance_strength'] = request.guidance_strength
+
+    texture_overrides = {}
+    if request.steps is not None:
+        texture_overrides['steps'] = request.steps
+    if request.texture_guidance is not None:
+        texture_overrides['guidance_strength'] = request.texture_guidance
+
+    # Generate mesh
+    try:
+        meshes = pipeline.run(
+            image,
+            seed=request.seed,
+            pipeline_type=request.pipeline_type,
+            sparse_structure_sampler_params=structure_shape_overrides,
+            shape_slat_sampler_params=structure_shape_overrides,
+            tex_slat_sampler_params=texture_overrides,
+        )
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        raise HTTPException(status_code=500, detail=f"Generation failed: {e}")
+
+    mesh = meshes[0]
+
+    # Export to GLB
+    try:
+        if request.output_path:
+            out_dir = os.path.dirname(request.output_path)
+            if out_dir:
+                os.makedirs(out_dir, exist_ok=True)
+            glb_path = request.output_path
+        else:
+            tmp = tempfile.NamedTemporaryFile(suffix=".glb", delete=False)
+            glb_path = tmp.name
+            tmp.close()
+
+        from mlx_backend.pipeline import to_glb
+        to_glb(
+            mesh, glb_path,
+            decimation_target=request.decimation_target,
+            texture_size=request.texture_size,
+            remesh=request.remesh,
+            verbose=True,
+        )
+        with open(glb_path, "rb") as f:
+            glb_bytes = f.read()
+        if not request.output_path:
+            os.unlink(glb_path)
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        raise HTTPException(status_code=500, detail=f"GLB export failed: {e}")
+
+    t_end = time.time()
+
+    return GenerateResponse(
+        glb=base64.b64encode(glb_bytes).decode(),
+        vertices=int(mesh.vertices.shape[0]),
+        faces=int(mesh.faces.shape[0]),
+        generation_time=round(t_end - t_start, 2),
+    )
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Trellis2 MLX API Server")
+    parser.add_argument("--host", type=str, default="0.0.0.0")
+    parser.add_argument("--port", type=int, default=8082)
+    parser.add_argument("--weights", type=str, default="weights/TRELLIS.2-4B")
+    args = parser.parse_args()
+
+    os.environ["TRELLIS2_WEIGHTS"] = args.weights
+    setup_logging()
+    uvicorn.run(app, host=args.host, port=args.port)
+
+
+if __name__ == "__main__":
+    main()

app_mlx.py

@@ -0,0 +1,174 @@
+"""
+Gradio app for Trellis2 with MLX backend (macOS).
+Simplified from upstream app.py — no CUDA render preview, direct GLB export.
+"""
+import gradio as gr
+import os
+import time
+from datetime import datetime
+import shutil
+import numpy as np
+from PIL import Image
+import torch
+import o_voxel
+
+MAX_SEED = np.iinfo(np.int32).max
+TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
+
+
+def get_seed(randomize_seed: bool, seed: int) -> int:
+    return np.random.randint(0, MAX_SEED) if randomize_seed else seed
+
+
+def preprocess_image(image: Image.Image) -> Image.Image:
+    return pipeline.preprocess_image(image)
+
+
+def image_to_3d(
+    image: Image.Image,
+    seed: int,
+    resolution: str,
+    ss_guidance_strength: float,
+    ss_sampling_steps: int,
+    shape_slat_guidance_strength: float,
+    shape_slat_sampling_steps: int,
+    tex_slat_guidance_strength: float,
+    tex_slat_sampling_steps: int,
+    decimation_target: int,
+    texture_size: int,
+    req: gr.Request,
+    progress=gr.Progress(track_tqdm=True),
+):
+    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    os.makedirs(user_dir, exist_ok=True)
+
+    pipeline_type = {"512": "512", "1024": "1024_cascade", "1536": "1536_cascade"}[resolution]
+
+    t0 = time.time()
+    meshes = pipeline.run(
+        image,
+        seed=seed,
+        sparse_structure_sampler_params={
+            "steps": ss_sampling_steps,
+            "guidance_strength": ss_guidance_strength,
+        },
+        shape_slat_sampler_params={
+            "steps": shape_slat_sampling_steps,
+            "guidance_strength": shape_slat_guidance_strength,
+        },
+        tex_slat_sampler_params={
+            "steps": tex_slat_sampling_steps,
+            "guidance_strength": tex_slat_guidance_strength,
+        },
+        pipeline_type=pipeline_type,
+    )
+    dt_gen = time.time() - t0
+    mesh = meshes[0]
+
+    t0 = time.time()
+    glb = o_voxel.postprocess.to_glb(
+        vertices=mesh.vertices,
+        faces=mesh.faces,
+        attr_volume=mesh.attrs,
+        coords=mesh.coords,
+        attr_layout=mesh.layout,
+        voxel_size=mesh.voxel_size,
+        aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
+        decimation_target=decimation_target,
+        texture_size=texture_size,
+        verbose=True,
+    )
+    dt_post = time.time() - t0
+
+    now = datetime.now()
+    timestamp = now.strftime("%Y-%m-%dT%H%M%S") + f".{now.microsecond // 1000:03d}"
+    glb_path = os.path.join(user_dir, f'sample_{timestamp}.glb')
+    glb.export(glb_path)
+
+    info = (f"Generation: {dt_gen:.0f}s | Post-processing: {dt_post:.0f}s | "
+            f"Verts: {mesh.vertices.shape[0]:,} | Faces: {mesh.faces.shape[0]:,}")
+    return glb_path, glb_path, info
+
+
+def start_session(req: gr.Request):
+    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    os.makedirs(user_dir, exist_ok=True)
+
+
+def end_session(req: gr.Request):
+    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    if os.path.exists(user_dir):
+        shutil.rmtree(user_dir)
+
+
+with gr.Blocks(title="Trellis2 MLX") as demo:
+    gr.Markdown("""
+    ## Trellis2 (MLX Backend)
+    Upload an image and click Generate to create a 3D asset.
+    """)
+
+    with gr.Row():
+        with gr.Column(scale=1, min_width=360):
+            image_prompt = gr.Image(label="Image Prompt", format="png", image_mode="RGBA", type="pil", height=400)
+
+            resolution = gr.Radio(["512", "1024"], label="Resolution", value="1024")
+            seed = gr.Slider(0, MAX_SEED, label="Seed", value=42, step=1)
+            randomize_seed = gr.Checkbox(label="Randomize Seed", value=False)
+            decimation_target = gr.Slider(100000, 1000000, label="Decimation Target", value=1000000, step=10000)
+            texture_size = gr.Slider(1024, 4096, label="Texture Size", value=2048, step=1024)
+
+            generate_btn = gr.Button("Generate", variant="primary")
+
+            with gr.Accordion(label="Advanced Settings", open=False):
+                gr.Markdown("### Stage 1: Sparse Structure")
+                with gr.Row():
+                    ss_guidance_strength = gr.Slider(1.0, 10.0, label="Guidance", value=7.5, step=0.1)
+                    ss_sampling_steps = gr.Slider(1, 50, label="Steps", value=12, step=1)
+                gr.Markdown("### Stage 2: Shape")
+                with gr.Row():
+                    shape_slat_guidance_strength = gr.Slider(1.0, 10.0, label="Guidance", value=7.5, step=0.1)
+                    shape_slat_sampling_steps = gr.Slider(1, 50, label="Steps", value=12, step=1)
+                gr.Markdown("### Stage 3: Texture")
+                with gr.Row():
+                    tex_slat_guidance_strength = gr.Slider(0.1, 10.0, label="Guidance", value=1.0, step=0.1)
+                    tex_slat_sampling_steps = gr.Slider(1, 50, label="Steps", value=12, step=1)
+
+        with gr.Column(scale=2):
+            glb_output = gr.Model3D(label="Generated 3D Model", height=700, clear_color=(0.25, 0.25, 0.25, 1.0))
+            download_btn = gr.DownloadButton(label="Download GLB")
+            info_text = gr.Textbox(label="Info", interactive=False)
+
+    # Handlers
+    demo.load(start_session)
+    demo.unload(end_session)
+
+    image_prompt.upload(
+        preprocess_image,
+        inputs=[image_prompt],
+        outputs=[image_prompt],
+    )
+
+    generate_btn.click(
+        get_seed,
+        inputs=[randomize_seed, seed],
+        outputs=[seed],
+    ).then(
+        image_to_3d,
+        inputs=[
+            image_prompt, seed, resolution,
+            ss_guidance_strength, ss_sampling_steps,
+            shape_slat_guidance_strength, shape_slat_sampling_steps,
+            tex_slat_guidance_strength, tex_slat_sampling_steps,
+            decimation_target, texture_size,
+        ],
+        outputs=[glb_output, download_btn, info_text],
+    )
+
+
+if __name__ == "__main__":
+    os.makedirs(TMP_DIR, exist_ok=True)
+
+    from mlx_backend.pipeline import create_mlx_pipeline
+    pipeline = create_mlx_pipeline(weights_path="weights/TRELLIS.2-4B")
+
+    demo.launch()

mlx_backend/__init__.py

@@ -0,0 +1,120 @@
+"""
+MLX backend for Trellis2 on Apple Silicon.
+
+All models run in pure MLX — no PyTorch except at the mesh-extraction boundary.
+Weight loading uses mx.load() for zero-copy safetensors reads.
+"""
+import logging
+import mlx.core as mx
+
+
+def setup_logging(level=logging.INFO):
+    """Enable MLX backend logging. Call before pipeline.run()."""
+    logging.basicConfig(
+        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
+        datefmt="%H:%M:%S",
+    )
+    logging.getLogger("mlx_backend").setLevel(level)
+
+__all__ = [
+    'load_safetensors',
+    'remap_flow_model_weights',
+    'remap_vae_decoder_weights',
+]
+
+
+def load_safetensors(path: str, dtype=None) -> dict:
+    """Load safetensors weights via MLX.
+
+    Keeps native dtype (bf16/fp16) by default. The RoPE interleaving fix
+    ensures numerical correctness regardless of weight precision.
+    Pass dtype=mx.float32 to force float32 if needed for debugging.
+    """
+    weights = mx.load(path)
+    if dtype is not None:
+        weights = {k: v.astype(dtype) if v.dtype in (mx.bfloat16, mx.float16) else v
+                   for k, v in weights.items()}
+    return weights
+
+
+def remap_flow_model_weights(weights: dict) -> dict:
+    """
+    Remap safetensors keys to MLX module paths for flow models.
+
+    MLX uses plain list indices for model.blocks (no `.layers.`):
+      blocks.0.xxx stays as blocks.0.xxx
+
+    But nn.Sequential-like containers use `.layers.`:
+      blocks.0.mlp.mlp.0.weight → blocks.0.mlp.mlp.layers.0.weight
+      adaLN_modulation.1.weight → adaLN_modulation.layers.1.weight
+      t_embedder.mlp.0.weight   → t_embedder.mlp.layers.0.weight
+    """
+    remapped = {}
+    for k, v in weights.items():
+        new_k = k
+        # blocks.N.xxx stays as-is (MLX uses integer indices for lists)
+        # Only remap Sequential containers within blocks
+
+        # Sequential .N. → .layers.N. for specific containers
+        new_k = _remap_sequential(new_k)
+
+        remapped[new_k] = v
+    return remapped
+
+
+def remap_vae_decoder_weights(weights: dict) -> dict:
+    """
+    Remap safetensors keys to MLX module paths for VAE decoders.
+
+    VAE decoder has nested lists: blocks[i][j].xxx
+    MLX uses plain integer indices: blocks.i.j.xxx (no .layers.)
+    """
+    remapped = {}
+    for k, v in weights.items():
+        new_k = k
+
+        # MlxSparseLinear wraps nn.Linear as self.linear
+        # from_latent.weight → from_latent.linear.weight
+        # output_layer.weight → output_layer.linear.weight
+        # Also handle to_subdiv in upsample blocks
+        for linear_name in ['from_latent', 'output_layer', 'to_subdiv', 'skip_connection']:
+            if new_k.startswith(f'{linear_name}.'):
+                new_k = new_k.replace(f'{linear_name}.', f'{linear_name}.linear.', 1)
+                break
+            # Handle nested: blocks.X.Y.to_subdiv.weight
+            import re
+            pattern = rf'(blocks\.\d+\.\d+\.{linear_name})\.'
+            new_k_try = re.sub(pattern, rf'\1.linear.', new_k, count=1)
+            if new_k_try != new_k:
+                new_k = new_k_try
+                break
+
+        # Sequential .N. → .layers.N. for MLP-like containers
+        new_k = _remap_sequential(new_k)
+
+        remapped[new_k] = v
+    return remapped
+
+
+def _remap_sequential(key: str) -> str:
+    """
+    Convert PyTorch nn.Sequential index notation to MLX.
+    e.g. 'foo.0.weight' → 'foo.layers.0.weight' when '0' is a digit
+    within known sequential containers.
+    """
+    import re
+    # Match patterns like: prefix.DIGIT.suffix where prefix ends with
+    # a known sequential container name
+    sequential_containers = [
+        'mlp.mlp', 'adaLN_modulation', 't_embedder.mlp', 'mlp'
+    ]
+    for container in sequential_containers:
+        # Pattern: container.DIGIT.rest
+        pattern = re.escape(container) + r'\.(\d+)\.'
+        replacement = container + r'.layers.\1.'
+        key = re.sub(pattern, replacement, key)
+        # Also handle terminal case: container.DIGIT (no trailing dot)
+        pattern_end = re.escape(container) + r'\.(\d+)$'
+        replacement_end = container + r'.layers.\1'
+        key = re.sub(pattern_end, replacement_end, key)
+    return key

mlx_backend/adapters.py

@@ -0,0 +1,264 @@
+"""
+Thin adapters wrapping MLX models to match upstream PyTorch model interfaces.
+Used by create_mlx_pipeline() so the upstream Trellis2ImageTo3DPipeline can
+call MLX models transparently via torch→numpy→mx→numpy→torch conversion.
+"""
+import numpy as np
+import torch
+import torch.nn as nn
+import mlx.core as mx
+
+from .sparse_tensor import MlxSparseTensor
+from trellis2.modules.sparse import SparseTensor
+
+
+# ---------------------------------------------------------------------------
+# Conversion helpers
+# ---------------------------------------------------------------------------
+
+def _sparse_to_mlx(st: SparseTensor) -> MlxSparseTensor:
+    """Convert upstream SparseTensor to MlxSparseTensor."""
+    feats = mx.array(st.feats.cpu().numpy())
+    coords = mx.array(st.coords.cpu().numpy().astype(np.int32))
+    return MlxSparseTensor(feats=feats, coords=coords)
+
+
+def _mlx_to_sparse(mx_st: MlxSparseTensor, ref: SparseTensor = None) -> SparseTensor:
+    """Convert MlxSparseTensor back to upstream SparseTensor.
+
+    If ref is provided and coords match, reuses backend data via .replace().
+    """
+    mx.eval(mx_st.feats)
+    feats = torch.from_numpy(np.array(mx_st.feats))
+    if ref is not None and mx_st.coords.shape == ref.coords.shape:
+        return ref.replace(feats=feats)
+    mx.eval(mx_st.coords)
+    coords = torch.from_numpy(np.array(mx_st.coords)).int()
+    return SparseTensor(feats=feats, coords=coords)
+
+
+def _torch_to_mx(t: torch.Tensor) -> mx.array:
+    """Convert a PyTorch CPU tensor to MLX array."""
+    if t.dtype == torch.bfloat16:
+        return mx.array(t.float().cpu().numpy()).astype(mx.bfloat16)
+    return mx.array(t.cpu().numpy())
+
+
+def _mx_to_torch(a: mx.array) -> torch.Tensor:
+    """Convert MLX array to PyTorch CPU tensor."""
+    mx.eval(a)
+    return torch.from_numpy(np.array(a))
+
+
+# ---------------------------------------------------------------------------
+# MlxFlowModelAdapter
+# ---------------------------------------------------------------------------
+
+class MlxFlowModelAdapter(nn.Module):
+    """Wraps dense MlxSparseStructureFlowModel or sparse MlxSLatFlowModel.
+
+    The upstream PT sampler calls model(x_t, t, cond, **kwargs).
+    This adapter converts tensors, runs the MLX model, and converts back.
+    """
+
+    def __init__(self, mlx_model, is_sparse: bool = False):
+        super().__init__()
+        self._mlx = mlx_model
+        self._is_sparse = is_sparse
+        self.resolution = mlx_model.resolution
+        self.in_channels = mlx_model.in_channels
+        self.out_channels = mlx_model.out_channels
+
+    def forward(self, x, t, cond, **kwargs):
+        if self._is_sparse:
+            return self._forward_sparse(x, t, cond, **kwargs)
+        return self._forward_dense(x, t, cond)
+
+    def _forward_dense(self, x, t, cond):
+        mx_out = self._mlx(_torch_to_mx(x), _torch_to_mx(t), _torch_to_mx(cond))
+        return _mx_to_torch(mx_out)
+
+    def _forward_sparse(self, x, t, cond, **kwargs):
+        mx_kwargs = {}
+        if kwargs.get('concat_cond') is not None:
+            mx_kwargs['concat_cond'] = _sparse_to_mlx(kwargs['concat_cond'])
+        mx_out = self._mlx(
+            _sparse_to_mlx(x), _torch_to_mx(t), _torch_to_mx(cond), **mx_kwargs
+        )
+        return _mlx_to_sparse(mx_out, x)
+
+    def to(self, *args, **kwargs):
+        return self
+
+    def cpu(self):
+        return self
+
+    def eval(self):
+        return self
+
+
+# ---------------------------------------------------------------------------
+# MlxStructureDecoderAdapter
+# ---------------------------------------------------------------------------
+
+class MlxStructureDecoderAdapter(nn.Module):
+    """Wraps MlxSparseStructureDecoder for the upstream pipeline."""
+
+    def __init__(self, mlx_model):
+        super().__init__()
+        self._mlx = mlx_model
+
+    def forward(self, z_s):
+        return _mx_to_torch(self._mlx(_torch_to_mx(z_s)))
+
+    def to(self, *args, **kwargs):
+        return self
+
+    def cpu(self):
+        return self
+
+    def eval(self):
+        return self
+
+
+# ---------------------------------------------------------------------------
+# MlxFlexiDualGridAdapter
+# ---------------------------------------------------------------------------
+
+class MlxFlexiDualGridAdapter(nn.Module):
+    """Wraps MlxFlexiDualGridVaeDecoder to match upstream FlexiDualGridVaeDecoder.
+
+    Upstream returns (meshes, subs) in eval mode with return_subs=True,
+    where meshes is a list of Mesh objects. This adapter runs the MLX decoder,
+    converts outputs to torch, then does mesh extraction via o_voxel.
+    """
+
+    def __init__(self, mlx_model):
+        super().__init__()
+        self._mlx = mlx_model
+        self.low_vram = False  # upstream toggles this
+
+    @property
+    def resolution(self):
+        return self._mlx.resolution
+
+    def set_resolution(self, resolution):
+        self._mlx.set_resolution(resolution)
+
+    def forward(self, x, return_subs=False, **kwargs):
+        from trellis2.representations import Mesh
+        from o_voxel.convert import flexible_dual_grid_to_mesh
+
+        mx_x = _sparse_to_mlx(x)
+        result = self._mlx(mx_x, return_subs=return_subs)
+
+        if return_subs:
+            (h_mx, verts_mx, inter_mx, quad_mx), subs_mx = result
+        else:
+            (h_mx, verts_mx, inter_mx, quad_mx) = result
+            subs_mx = None
+
+        mx.eval(h_mx.feats, h_mx.coords, verts_mx, inter_mx, quad_mx)
+
+        # Convert to torch for mesh extraction
+        coords_t = torch.from_numpy(np.array(h_mx.coords[:, 1:])).int()
+        verts_t = torch.from_numpy(np.array(verts_mx)).float()
+        inter_t = torch.from_numpy(np.array(inter_mx)).bool()
+        quad_t = torch.from_numpy(np.array(quad_mx)).float()
+
+        mesh_verts, mesh_faces = flexible_dual_grid_to_mesh(
+            coords_t, verts_t, inter_t, quad_t,
+            aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
+            grid_size=self._mlx.resolution,
+            train=False,
+        )
+        meshes = [Mesh(mesh_verts, mesh_faces)]
+
+        if return_subs:
+            subs_t = []
+            for s in subs_mx:
+                mx.eval(s.feats)
+                subs_t.append(_mlx_to_sparse(s))
+            return meshes, subs_t
+
+        return meshes
+
+    def upsample(self, x, upsample_times):
+        mx_x = _sparse_to_mlx(x)
+        mx_coords = self._mlx.upsample(mx_x, upsample_times)
+        mx.eval(mx_coords)
+        return torch.from_numpy(np.array(mx_coords)).int()
+
+    def to(self, *args, **kwargs):
+        return self
+
+    def cpu(self):
+        return self
+
+    def eval(self):
+        return self
+
+
+# ---------------------------------------------------------------------------
+# MlxTexVaeDecoderAdapter
+# ---------------------------------------------------------------------------
+
+class MlxTexVaeDecoderAdapter(nn.Module):
+    """Wraps MlxSparseUnetVaeDecoder (texture) for the upstream pipeline.
+
+    Returns SparseTensor so upstream arithmetic (* 0.5 + 0.5) works.
+    """
+
+    def __init__(self, mlx_model):
+        super().__init__()
+        self._mlx = mlx_model
+
+    def forward(self, x, guide_subs=None, **kwargs):
+        mx_x = _sparse_to_mlx(x)
+        mx_guide = None
+        if guide_subs is not None:
+            mx_guide = [_sparse_to_mlx(s) for s in guide_subs]
+
+        result = self._mlx(mx_x, guide_subs=mx_guide)
+        mx.eval(result.feats)
+        return _mlx_to_sparse(result, x)
+
+    def to(self, *args, **kwargs):
+        return self
+
+    def cpu(self):
+        return self
+
+    def eval(self):
+        return self
+
+
+# ---------------------------------------------------------------------------
+# MlxImageCondAdapter
+# ---------------------------------------------------------------------------
+
+class MlxImageCondAdapter:
+    """Wraps MlxDINOv3FeatureExtractor for the upstream pipeline.
+
+    The upstream pipeline sets .image_size then calls model(images).
+    """
+
+    def __init__(self, mlx_dino):
+        self._mlx = mlx_dino
+        self.image_size = 512
+
+    def __call__(self, images):
+        from PIL import Image as PILImage
+        resized = [
+            img.resize((self.image_size, self.image_size), PILImage.LANCZOS)
+            for img in images
+        ]
+        mx_out = self._mlx(resized)
+        mx.eval(mx_out)
+        return torch.from_numpy(np.array(mx_out))
+
+    def to(self, *args, **kwargs):
+        return self
+
+    def cpu(self):
+        return self

mlx_backend/attention.py

@@ -0,0 +1,145 @@
+"""
+Multi-head attention for MLX backend.
+Fused scaled dot-product attention via Metal kernel, with QK-RMSNorm and RoPE.
+
+Supports two modes:
+- Sparse (unbatched): x is (N, C) — used for sparse flow models
+- Dense (batched): x is (B, N, C) — used for dense structure flow with batched CFG
+"""
+import mlx.core as mx
+import mlx.nn as nn
+from .norm import SparseMultiHeadRMSNorm
+from .rope import build_rope_freqs, compute_rope_phases, apply_rope
+
+
+class MlxMultiHeadAttention(nn.Module):
+    """
+    Multi-head attention matching the PyTorch SparseMultiHeadAttention.
+
+    For self-attention: fused QKV projection, optional QK-RMSNorm, optional RoPE.
+    For cross-attention: separate Q and KV projections.
+
+    Supports both (N, C) unbatched and (B, N, C) batched input.
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        num_heads: int,
+        ctx_channels: int = None,
+        type: str = "self",
+        qkv_bias: bool = True,
+        use_rope: bool = False,
+        rope_freq: tuple = (1.0, 10000.0),
+        qk_rms_norm: bool = False,
+    ):
+        super().__init__()
+        assert channels % num_heads == 0
+        self.channels = channels
+        self.num_heads = num_heads
+        self.head_dim = channels // num_heads
+        self.ctx_channels = ctx_channels or channels
+        self._type = type
+        self.use_rope = use_rope
+        self.qk_rms_norm = qk_rms_norm
+
+        if type == "self":
+            self.to_qkv = nn.Linear(channels, channels * 3, bias=qkv_bias)
+        else:
+            self.to_q = nn.Linear(channels, channels, bias=qkv_bias)
+            self.to_kv = nn.Linear(self.ctx_channels, channels * 2, bias=qkv_bias)
+
+        if qk_rms_norm:
+            self.q_rms_norm = SparseMultiHeadRMSNorm(self.head_dim, num_heads)
+            self.k_rms_norm = SparseMultiHeadRMSNorm(self.head_dim, num_heads)
+
+        self.to_out = nn.Linear(channels, channels)
+
+        if use_rope:
+            self._rope_freqs = build_rope_freqs(self.head_dim, dim=3, rope_freq=rope_freq)
+
+    def __call__(
+        self,
+        x: mx.array,
+        context: mx.array = None,
+        rope_cache: tuple = None,
+    ) -> mx.array:
+        """
+        Args:
+            x: (N, C) or (B, N, C) input features
+            context: (M, ctx_C) or (B, M, ctx_C) cross-attention context
+            rope_cache: (cos, sin) precomputed RoPE phases
+
+        Returns:
+            Same shape as x
+        """
+        H = self.num_heads
+        D = self.head_dim
+        batched = x.ndim == 3
+
+        if self._type == "self":
+            qkv = self.to_qkv(x)  # (..., 3*C)
+            if batched:
+                B, N, _ = qkv.shape
+                qkv = qkv.reshape(B, N, 3, H, D)
+                q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]  # (B, N, H, D)
+            else:
+                qkv = qkv.reshape(-1, 3, H, D)
+                q, k, v = qkv[:, 0], qkv[:, 1], qkv[:, 2]  # (N, H, D)
+
+            if self.qk_rms_norm:
+                q = self.q_rms_norm(q)
+                k = self.k_rms_norm(k)
+
+            if self.use_rope and rope_cache is not None:
+                cos, sin = rope_cache
+                q = apply_rope(q, cos, sin)
+                k = apply_rope(k, cos, sin)
+
+            out = self._sdpa(q, k, v, batched)
+        else:
+            q = self.to_q(x)
+            kv = self.to_kv(context)
+            if batched:
+                B, N, _ = q.shape
+                q = q.reshape(B, N, H, D)
+                M = kv.shape[1]
+                kv = kv.reshape(B, M, 2, H, D)
+                k, v = kv[:, :, 0], kv[:, :, 1]  # (B, M, H, D)
+            else:
+                q = q.reshape(-1, H, D)
+                kv = kv.reshape(-1, 2, H, D)
+                k, v = kv[:, 0], kv[:, 1]
+
+            if self.qk_rms_norm:
+                q = self.q_rms_norm(q)
+                k = self.k_rms_norm(k)
+
+            out = self._sdpa(q, k, v, batched)
+
+        if batched:
+            out = out.reshape(B, -1, self.channels)
+        else:
+            out = out.reshape(-1, self.channels)
+        return self.to_out(out)
+
+    def _sdpa(self, q: mx.array, k: mx.array, v: mx.array, batched: bool = False) -> mx.array:
+        """
+        Fused scaled dot-product attention via Metal kernel.
+        Inputs: (N, H, D) unbatched or (B, N, H, D) batched.
+        """
+        scale = self.head_dim ** -0.5
+        if batched:
+            # (B, N, H, D) -> (B, H, N, D)
+            q = q.transpose(0, 2, 1, 3)
+            k = k.transpose(0, 2, 1, 3)
+            v = v.transpose(0, 2, 1, 3)
+            out = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale)
+            return out.transpose(0, 2, 1, 3)  # (B, N, H, D)
+        else:
+            # (N, H, D) -> (1, H, N, D)
+            q = q.transpose(1, 0, 2)[None]
+            k = k.transpose(1, 0, 2)[None]
+            v = v.transpose(1, 0, 2)[None]
+            out = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale)
+            return out[0].transpose(1, 0, 2)  # (N, H, D)

mlx_backend/convert.py

@@ -0,0 +1,26 @@
+"""
+Conversion utilities between PyTorch and MLX tensors.
+Used ONLY at pipeline boundaries (mesh extraction input).
+"""
+import numpy as np
+import mlx.core as mx
+
+
+def torch_to_mlx(t) -> mx.array:
+    """Convert a PyTorch tensor to MLX array."""
+    import torch
+    if t.dtype == torch.bfloat16:
+        return mx.array(t.float().cpu().numpy()).astype(mx.bfloat16)
+    return mx.array(t.cpu().numpy())
+
+
+def mlx_to_torch(a: mx.array):
+    """Convert an MLX array to PyTorch tensor."""
+    import torch
+    arr = np.array(a)
+    return torch.from_numpy(arr)
+
+
+def mlx_to_numpy(a: mx.array) -> np.ndarray:
+    """Convert an MLX array to numpy."""
+    return np.array(a)

mlx_backend/dinov3.py

@@ -0,0 +1,337 @@
+"""
+DINOv3 ViT feature extractor in MLX.
+Uses HuggingFace transformers weights but runs purely in MLX.
+"""
+import math
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+from PIL import Image
+
+
+class MlxDINOv3PatchEmbed(nn.Module):
+    """Patch embedding: Conv2d(3, dim, kernel=16, stride=16) + CLS + registers."""
+
+    def __init__(self, dim: int = 1024, patch_size: int = 16, num_register_tokens: int = 4):
+        super().__init__()
+        self.patch_size = patch_size
+        self.num_register_tokens = num_register_tokens
+        self.weight = mx.zeros((dim, patch_size, patch_size, 3))
+        self.bias = mx.zeros((dim,))
+        self.cls_token = mx.zeros((1, 1, dim))
+        self.register_tokens = mx.zeros((1, num_register_tokens, dim))
+
+    def __call__(self, x: mx.array) -> mx.array:
+        """x: (B, H, W, 3) -> (B, num_patches + 1 + num_reg, dim)"""
+        B, H, W, C = x.shape
+        P = self.patch_size
+        nH, nW = H // P, W // P
+
+        # Manual convolution via reshape + matmul (stride == kernel_size)
+        x = x.reshape(B, nH, P, nW, P, C)
+        x = x.transpose(0, 1, 3, 2, 4, 5)  # (B, nH, nW, P, P, C)
+        x = x.reshape(B, nH * nW, P * P * C)
+
+        w = self.weight.reshape(self.weight.shape[0], -1).T  # (P*P*3, dim)
+        patches = x @ w + self.bias  # (B, N, dim)
+
+        cls = mx.broadcast_to(self.cls_token, (B, 1, patches.shape[-1]))
+        reg = mx.broadcast_to(self.register_tokens, (B, self.num_register_tokens, patches.shape[-1]))
+        return mx.concatenate([cls, reg, patches], axis=1)
+
+
+class MlxDINOv3TransformerBlock(nn.Module):
+    """ViT transformer block with layer scaling."""
+
+    def __init__(self, dim: int = 1024, num_heads: int = 16, mlp_ratio: float = 4.0):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+        self.attn = MlxDINOv3Attention(dim, num_heads)
+        hidden = int(dim * mlp_ratio)
+        self.mlp = MlxDINOv3MLP(dim, hidden)
+        # Layer scaling
+        self.layer_scale1 = mx.ones((dim,))
+        self.layer_scale2 = mx.ones((dim,))
+
+    def __call__(self, x: mx.array, rope_cos: mx.array = None, rope_sin: mx.array = None,
+                 num_prefix_tokens: int = 5) -> mx.array:
+        h = self.norm1(x)
+        h = self.attn(h, rope_cos, rope_sin, num_prefix_tokens=num_prefix_tokens)
+        x = x + self.layer_scale1 * h
+        h = self.norm2(x)
+        h = self.mlp(h)
+        x = x + self.layer_scale2 * h
+        return x
+
+
+class MlxDINOv3Attention(nn.Module):
+    """Multi-head self-attention with RoPE (matching HF DINOv3ViTAttention).
+
+    Uses separate Q, K, V projections (not fused QKV) to match HF weight layout.
+    RoPE is applied only to patch tokens, skipping prefix (CLS + register) tokens.
+    """
+
+    def __init__(self, dim: int, num_heads: int):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.q_proj = nn.Linear(dim, dim, bias=True)
+        self.k_proj = nn.Linear(dim, dim, bias=False)  # DINOv3: key has no bias
+        self.v_proj = nn.Linear(dim, dim, bias=True)
+        self.o_proj = nn.Linear(dim, dim, bias=True)
+
+    def __call__(self, x: mx.array, rope_cos: mx.array = None, rope_sin: mx.array = None,
+                 num_prefix_tokens: int = 5) -> mx.array:
+        B, N, C = x.shape
+        H = self.num_heads
+        D = self.head_dim
+
+        q = self.q_proj(x).reshape(B, N, H, D).transpose(0, 2, 1, 3)  # (B, H, N, D)
+        k = self.k_proj(x).reshape(B, N, H, D).transpose(0, 2, 1, 3)
+        v = self.v_proj(x).reshape(B, N, H, D).transpose(0, 2, 1, 3)
+
+        if rope_cos is not None:
+            # Apply RoPE only to patch tokens, skip prefix (CLS + registers)
+            num_patches = rope_cos.shape[-2]
+            q_prefix = q[:, :, :num_prefix_tokens]
+            k_prefix = k[:, :, :num_prefix_tokens]
+            q_patches = q[:, :, num_prefix_tokens:]
+            k_patches = k[:, :, num_prefix_tokens:]
+
+            q_patches = self._apply_rope(q_patches, rope_cos, rope_sin)
+            k_patches = self._apply_rope(k_patches, rope_cos, rope_sin)
+
+            q = mx.concatenate([q_prefix, q_patches], axis=2)
+            k = mx.concatenate([k_prefix, k_patches], axis=2)
+
+        scale = D ** -0.5
+        out = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale)
+        out = out.transpose(0, 2, 1, 3).reshape(B, N, C)
+        return self.o_proj(out)
+
+    @staticmethod
+    def _apply_rope(x: mx.array, cos: mx.array, sin: mx.array) -> mx.array:
+        """Apply RoPE via rotate_half: x*cos + rotate_half(x)*sin.
+
+        cos/sin: (1, N_patches, head_dim) — full head_dim, already tiled.
+        x: (B, H, N_patches, head_dim).
+        """
+        # rotate_half: [-x2, x1]
+        half = x.shape[-1] // 2
+        x1 = x[..., :half]
+        x2 = x[..., half:]
+        x_rot = mx.concatenate([-x2, x1], axis=-1)
+        return x * cos + x_rot * sin
+
+
+class MlxDINOv3MLP(nn.Module):
+    """MLP with GELU activation."""
+
+    def __init__(self, dim: int, hidden: int):
+        super().__init__()
+        self.fc1 = nn.Linear(dim, hidden, bias=True)
+        self.fc2 = nn.Linear(hidden, dim, bias=True)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.fc2(nn.gelu(self.fc1(x)))
+
+
+class MlxDINOv3FeatureExtractor(nn.Module):
+    """
+    DINOv3 ViT-L/16 feature extractor in pure MLX.
+    24 layers, dim=1024, 16 heads, patch_size=16.
+    """
+
+    def __init__(self, dim: int = 1024, num_heads: int = 16, num_layers: int = 24,
+                 patch_size: int = 16, mlp_ratio: float = 4.0,
+                 num_register_tokens: int = 4, rope_theta: float = 100.0):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.patch_size = patch_size
+        self.num_register_tokens = num_register_tokens
+        self.rope_theta = rope_theta
+
+        self.embeddings = MlxDINOv3PatchEmbed(dim, patch_size, num_register_tokens)
+        self.layers = [MlxDINOv3TransformerBlock(dim, num_heads, mlp_ratio) for _ in range(num_layers)]
+        self.norm = nn.LayerNorm(dim)
+
+    def _build_rope_2d(self, h: int, w: int) -> tuple:
+        """Build 2D RoPE cos/sin matching HF DINOv3ViTRopePositionEmbedding.
+
+        Returns cos/sin for PATCH tokens only (prefix tokens get no RoPE).
+        Shape: (1, h*w, head_dim) — full head_dim, tiled 2x.
+        """
+        head_dim = self.dim // self.num_heads
+
+        # inv_freq: 1/theta^(arange(0, 1, 4/head_dim)) — matches HF exactly
+        inv_freq = 1.0 / (self.rope_theta ** mx.arange(0, 1, 4 / head_dim, dtype=mx.float32))
+        # inv_freq shape: (head_dim/4,)
+
+        # Patch center coords normalized to [-1, +1] — matches HF get_patches_center_coordinates
+        coords_h = (mx.arange(h, dtype=mx.float32) + 0.5) / h
+        coords_w = (mx.arange(w, dtype=mx.float32) + 0.5) / w
+        grid_h, grid_w = mx.meshgrid(coords_h, coords_w, indexing='ij')
+        coords = mx.stack([grid_h.reshape(-1), grid_w.reshape(-1)], axis=-1)  # (h*w, 2)
+        coords = 2.0 * coords - 1.0  # shift to [-1, +1]
+
+        # angles: 2π * coord * inv_freq, then flatten and tile 2x
+        angles = 2 * math.pi * coords[:, :, None] * inv_freq[None, None, :]  # (h*w, 2, head_dim/4)
+        angles = angles.reshape(h * w, -1)  # (h*w, head_dim/2)
+        # Tile to full head_dim (matching PT's angles.tile(2))
+        angles = mx.concatenate([angles, angles], axis=-1)  # (h*w, head_dim)
+
+        cos = mx.cos(angles)[None, :, :]  # (1, h*w, head_dim)
+        sin = mx.sin(angles)[None, :, :]
+        return cos, sin
+
+    def __call__(self, images: list) -> mx.array:
+        x = self._preprocess(images)
+        B, H, W, _ = x.shape
+
+        h = self.embeddings(x)
+
+        nH, nW = H // self.patch_size, W // self.patch_size
+        rope_cos, rope_sin = self._build_rope_2d(nH, nW)
+
+        # 1 CLS + num_register_tokens registers = prefix tokens (no RoPE)
+        num_prefix = 1 + self.num_register_tokens
+
+        for layer in self.layers:
+            h = layer(h, rope_cos, rope_sin, num_prefix_tokens=num_prefix)
+
+        # Match PT reference: pure layer_norm without learned params
+        # (PT's extract_features uses F.layer_norm, not the model's learned norm)
+        h = mx.fast.layer_norm(h, weight=None, bias=None, eps=1e-5)
+        return h
+
+    def _preprocess(self, images: list) -> mx.array:
+        """Resize, normalize, convert to MLX array."""
+        mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
+        std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
+
+        processed = []
+        for img in images:
+            if isinstance(img, Image.Image):
+                size = max(img.size)
+                target = ((size + self.patch_size - 1) // self.patch_size) * self.patch_size
+                img = img.resize((target, target), Image.LANCZOS)
+                arr = np.array(img.convert('RGB')).astype(np.float32) / 255.0
+            else:
+                arr = np.array(img, dtype=np.float32)
+            arr = (arr - mean) / std
+            processed.append(arr)
+
+        return mx.array(np.stack(processed))
+
+
+def load_dinov3_from_hf(model_name: str = "facebook/dinov3-vitl16-pretrain-lvd1689m",
+                        image_size: int = 512) -> MlxDINOv3FeatureExtractor:
+    """Load DINOv3 weights from HuggingFace into MLX model."""
+    from huggingface_hub import hf_hub_download
+    import json
+
+    config_path = hf_hub_download(model_name, "config.json")
+    weight_path = hf_hub_download(model_name, "model.safetensors")
+
+    with open(config_path) as f:
+        config = json.load(f)
+
+    dim = config.get('hidden_size', 1024)
+    num_heads = config.get('num_attention_heads', 16)
+    num_layers = config.get('num_hidden_layers', 24)
+    patch_size = config.get('patch_size', 16)
+    mlp_ratio = config.get('mlp_ratio', 4.0)
+    num_register = config.get('num_register_tokens', 4)
+
+    model = MlxDINOv3FeatureExtractor(
+        dim=dim, num_heads=num_heads, num_layers=num_layers,
+        patch_size=patch_size, mlp_ratio=mlp_ratio,
+        num_register_tokens=num_register,
+    )
+
+    weights = mx.load(weight_path)
+    remapped = _remap_dinov3_weights(weights, num_layers, dim)
+    model.load_weights(list(remapped.items()))
+
+    return model
+
+
+def _remap_dinov3_weights(weights: dict, num_layers: int, dim: int) -> dict:
+    """
+    Map HuggingFace DINOv3 weight keys to our model structure.
+
+    HF format:
+        layer.N.attention.q_proj.weight, layer.N.attention.k_proj.weight (no bias),
+        layer.N.attention.v_proj.weight, layer.N.attention.o_proj.weight/bias,
+        layer.N.mlp.up_proj.weight/bias, layer.N.mlp.down_proj.weight/bias,
+        layer.N.norm1/norm2.weight/bias, layer.N.layer_scale1/2.lambda1
+
+    Our format:
+        layers.N.attn.qkv.weight/bias, layers.N.attn.proj.weight/bias,
+        layers.N.mlp.fc1.weight/bias, layers.N.mlp.fc2.weight/bias,
+        layers.N.norm1/norm2.weight/bias, layers.N.layer_scale1/2
+    """
+    remapped = {}
+
+    # Embeddings
+    if 'embeddings.patch_embeddings.weight' in weights:
+        # HF: (dim, 3, P, P) -> our: (dim, P, P, 3)
+        remapped['embeddings.weight'] = weights['embeddings.patch_embeddings.weight'].transpose(0, 2, 3, 1)
+    if 'embeddings.patch_embeddings.bias' in weights:
+        remapped['embeddings.bias'] = weights['embeddings.patch_embeddings.bias']
+    if 'embeddings.cls_token' in weights:
+        remapped['embeddings.cls_token'] = weights['embeddings.cls_token']
+    if 'embeddings.register_tokens' in weights:
+        remapped['embeddings.register_tokens'] = weights['embeddings.register_tokens']
+
+    # Final norm
+    if 'norm.weight' in weights:
+        remapped['norm.weight'] = weights['norm.weight']
+    if 'norm.bias' in weights:
+        remapped['norm.bias'] = weights['norm.bias']
+
+    # Transformer layers — separate Q/K/V projections (matching HF layout)
+    for i in range(num_layers):
+        prefix_hf = f'layer.{i}'
+        prefix_mlx = f'layers.{i}'
+
+        # Q, K, V projections — keep separate (not fused)
+        for proj in ['q_proj', 'k_proj', 'v_proj']:
+            for suffix in ['weight', 'bias']:
+                key = f'{prefix_hf}.attention.{proj}.{suffix}'
+                if key in weights:
+                    remapped[f'{prefix_mlx}.attn.{proj}.{suffix}'] = weights[key]
+
+        # Output projection
+        for suffix in ['weight', 'bias']:
+            key = f'{prefix_hf}.attention.o_proj.{suffix}'
+            if key in weights:
+                remapped[f'{prefix_mlx}.attn.o_proj.{suffix}'] = weights[key]
+
+        # Norms
+        for norm in ['norm1', 'norm2']:
+            for suffix in ['weight', 'bias']:
+                key = f'{prefix_hf}.{norm}.{suffix}'
+                if key in weights:
+                    remapped[f'{prefix_mlx}.{norm}.{suffix}'] = weights[key]
+
+        # MLP: up_proj → fc1, down_proj → fc2
+        for suffix in ['weight', 'bias']:
+            key = f'{prefix_hf}.mlp.up_proj.{suffix}'
+            if key in weights:
+                remapped[f'{prefix_mlx}.mlp.fc1.{suffix}'] = weights[key]
+            key = f'{prefix_hf}.mlp.down_proj.{suffix}'
+            if key in weights:
+                remapped[f'{prefix_mlx}.mlp.fc2.{suffix}'] = weights[key]
+
+        # Layer scaling
+        key1 = f'{prefix_hf}.layer_scale1.lambda1'
+        if key1 in weights:
+            remapped[f'{prefix_mlx}.layer_scale1'] = weights[key1]
+        key2 = f'{prefix_hf}.layer_scale2.lambda1'
+        if key2 in weights:
+            remapped[f'{prefix_mlx}.layer_scale2'] = weights[key2]
+
+    return remapped

mlx_backend/flow_models.py

@@ -0,0 +1,353 @@
+"""
+Flow matching models in MLX.
+MlxSparseStructureFlowModel: dense 16^3 structure flow
+MlxSLatFlowModel: sparse structured latent flow (shape + texture)
+"""
+import logging
+import math
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+from .norm import LayerNorm32
+from .transformer_block import (
+    MlxModulatedSparseTransformerCrossBlock,
+    MlxModulatedTransformerCrossBlock,
+    MlxSparseSequential,
+)
+from .rope import build_rope_freqs, compute_rope_phases
+from .sparse_tensor import MlxSparseTensor, mlx_sparse_cat
+
+logger = logging.getLogger(__name__)
+
+
+def _metal_mem_mb() -> str:
+    """Current Metal memory usage in MB."""
+    try:
+        active = mx.get_active_memory() / 1024**2
+        peak = mx.get_peak_memory() / 1024**2
+        return f"{active:.0f}MB (peak {peak:.0f}MB)"
+    except Exception:
+        return "N/A"
+
+
+class MlxTimestepEmbedder(nn.Module):
+    """Sinusoidal timestep embedding → MLP."""
+
+    def __init__(self, hidden_size: int, frequency_embedding_size: int = 256):
+        super().__init__()
+        self.frequency_embedding_size = frequency_embedding_size
+        self.mlp = MlxSparseSequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+
+    def __call__(self, t: mx.array) -> mx.array:
+        t_freq = self._timestep_embedding(t, self.frequency_embedding_size)
+        return self.mlp(t_freq)
+
+    @staticmethod
+    def _timestep_embedding(t: mx.array, dim: int, max_period: float = 10000.0) -> mx.array:
+        half = dim // 2
+        freqs = mx.exp(
+            -math.log(max_period) * mx.arange(half, dtype=mx.float32) / half
+        )
+        args = t[:, None].astype(mx.float32) * freqs[None]
+        embedding = mx.concatenate([mx.cos(args), mx.sin(args)], axis=-1)
+        if dim % 2:
+            embedding = mx.concatenate([embedding, mx.zeros_like(embedding[:, :1])], axis=-1)
+        return embedding
+
+
+class MlxSparseStructureFlowModel(nn.Module):
+    """
+    Dense flow model for sparse structure sampling.
+    Input: (B, 8, 16, 16, 16) noise → (B, 8, 16, 16, 16) output.
+    Internally flattened to (B, 4096, C) dense sequence.
+    """
+
+    def __init__(
+        self,
+        resolution: int = 16,
+        in_channels: int = 8,
+        model_channels: int = 1536,
+        cond_channels: int = 1024,
+        out_channels: int = 8,
+        num_blocks: int = 30,
+        num_heads: int = 12,
+        mlp_ratio: float = 5.3334,
+        pe_mode: str = "rope",
+        share_mod: bool = True,
+        qk_rms_norm: bool = True,
+        qk_rms_norm_cross: bool = True,
+    ):
+        super().__init__()
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.model_channels = model_channels
+        self.num_heads = num_heads
+        self.pe_mode = pe_mode
+        self.share_mod = share_mod
+
+        self.t_embedder = MlxTimestepEmbedder(model_channels)
+        if share_mod:
+            self.adaLN_modulation = MlxSparseSequential(
+                nn.SiLU(),
+                nn.Linear(model_channels, 6 * model_channels, bias=True),
+            )
+
+        self.input_layer = nn.Linear(in_channels, model_channels)
+
+        self.blocks = [
+            MlxModulatedTransformerCrossBlock(
+                model_channels, cond_channels,
+                num_heads=num_heads, mlp_ratio=mlp_ratio,
+                use_rope=(pe_mode == "rope"),
+                share_mod=share_mod,
+                qk_rms_norm=qk_rms_norm,
+                qk_rms_norm_cross=qk_rms_norm_cross,
+            )
+            for _ in range(num_blocks)
+        ]
+
+        self.out_layer = nn.Linear(model_channels, out_channels)
+
+        # Precompute RoPE for 16^3 grid
+        if pe_mode == "rope":
+            head_dim = model_channels // num_heads
+            freqs = build_rope_freqs(head_dim, dim=3)
+            coords_np = np.stack(np.meshgrid(
+                np.arange(resolution), np.arange(resolution), np.arange(resolution),
+                indexing='ij'
+            ), axis=-1).reshape(-1, 3)
+            coords_mx = mx.array(coords_np.astype(np.int32))
+            self._rope_cos, self._rope_sin = compute_rope_phases(coords_mx, freqs, head_dim)
+
+        self._compiled_blocks = None
+
+    def _run_blocks(self, h, t_emb, cond, rope_cos, rope_sin):
+        rope_cache = (rope_cos, rope_sin)
+        for block in self.blocks:
+            h = block(h, t_emb, cond, rope_cache=rope_cache)
+        return h
+
+    def _run_blocks_no_rope(self, h, t_emb, cond):
+        for block in self.blocks:
+            h = block(h, t_emb, cond, rope_cache=None)
+        return h
+
+    def __call__(self, x: mx.array, t: mx.array, cond: mx.array) -> mx.array:
+        B = x.shape[0]
+        R = self.resolution
+
+        logger.debug("[MLX] StructureFlow: x=%s dtype=%s, mem=%s",
+                     x.shape, x.dtype, _metal_mem_mb())
+
+        # Flatten: (B, C, D, H, W) -> (B, D*H*W, C)
+        h = x.reshape(B, self.in_channels, -1).transpose(0, 2, 1)
+        h = self.input_layer(h)
+
+        # Match upstream bfloat16 reduced-precision casting (manual_cast)
+        compute_dtype = mx.bfloat16
+        h = h.astype(compute_dtype)
+
+        t_emb = self.t_embedder(t)
+        if self.share_mod:
+            t_emb = self.adaLN_modulation(t_emb)
+        t_emb = t_emb.astype(compute_dtype)
+        cond = cond.astype(compute_dtype)
+
+        if self.pe_mode == "rope":
+            if self._compiled_blocks is None:
+                try:
+                    self._compiled_blocks = mx.compile(self._run_blocks)
+                    logger.info("[MLX] StructureFlow: using mx.compile for block loop")
+                except Exception as e:
+                    logger.warning("[MLX] StructureFlow: mx.compile failed (%s), falling back to per-block eval", e)
+                    self._compiled_blocks = False
+            if self._compiled_blocks:
+                h = self._compiled_blocks(h, t_emb, cond, self._rope_cos, self._rope_sin)
+            else:
+                rope_cache = (self._rope_cos, self._rope_sin)
+                for i, block in enumerate(self.blocks):
+                    h = block(h, t_emb, cond, rope_cache=rope_cache)
+                    if (i + 1) % 10 == 0:
+                        mx.eval(h)  # periodic eval to bound memory in fallback path
+        else:
+            if self._compiled_blocks is None:
+                try:
+                    self._compiled_blocks = mx.compile(self._run_blocks_no_rope)
+                    logger.info("[MLX] StructureFlow: using mx.compile for block loop (no rope)")
+                except Exception as e:
+                    logger.warning("[MLX] StructureFlow: mx.compile failed (%s), falling back to per-block eval", e)
+                    self._compiled_blocks = False
+            if self._compiled_blocks:
+                h = self._compiled_blocks(h, t_emb, cond)
+            else:
+                for i, block in enumerate(self.blocks):
+                    h = block(h, t_emb, cond, rope_cache=None)
+                    if (i + 1) % 10 == 0:
+                        mx.eval(h)
+
+        logger.debug("[MLX] StructureFlow blocks done, mem=%s", _metal_mem_mb())
+
+        # Cast back to float32 before final norm (matches upstream cast back to input dtype)
+        h = h.astype(mx.float32)
+        # Two-pass LayerNorm (matches LayerNorm32 / upstream F.layer_norm precision)
+        mean = mx.mean(h, axis=-1, keepdims=True)
+        h = h - mean
+        var = mx.mean(h * h, axis=-1, keepdims=True)
+        h = h * mx.rsqrt(var + 1e-5)
+        h = self.out_layer(h)
+
+        h = h.transpose(0, 2, 1).reshape(B, self.out_channels, R, R, R)
+        return h
+
+
+class MlxSLatFlowModel(nn.Module):
+    """
+    Sparse structured latent flow model.
+    Input: MlxSparseTensor with (N, in_channels) features.
+    """
+
+    def __init__(
+        self,
+        resolution: int = 64,
+        in_channels: int = 32,
+        model_channels: int = 1536,
+        cond_channels: int = 1024,
+        out_channels: int = 32,
+        num_blocks: int = 30,
+        num_heads: int = 12,
+        mlp_ratio: float = 5.3334,
+        pe_mode: str = "rope",
+        share_mod: bool = True,
+        qk_rms_norm: bool = True,
+        qk_rms_norm_cross: bool = True,
+    ):
+        super().__init__()
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.model_channels = model_channels
+        self.num_heads = num_heads
+        self.pe_mode = pe_mode
+        self.share_mod = share_mod
+
+        self.t_embedder = MlxTimestepEmbedder(model_channels)
+        if share_mod:
+            self.adaLN_modulation = MlxSparseSequential(
+                nn.SiLU(),
+                nn.Linear(model_channels, 6 * model_channels, bias=True),
+            )
+
+        self.input_layer = nn.Linear(in_channels, model_channels)
+
+        self.blocks = [
+            MlxModulatedSparseTransformerCrossBlock(
+                model_channels, cond_channels,
+                num_heads=num_heads, mlp_ratio=mlp_ratio,
+                use_rope=(pe_mode == "rope"),
+                share_mod=share_mod,
+                qk_rms_norm=qk_rms_norm,
+                qk_rms_norm_cross=qk_rms_norm_cross,
+            )
+            for _ in range(num_blocks)
+        ]
+
+        self.out_layer = nn.Linear(model_channels, out_channels)
+
+        if pe_mode == "rope":
+            head_dim = model_channels // num_heads
+            self._rope_freqs = build_rope_freqs(head_dim, dim=3)
+
+        self._compiled_blocks = None
+
+    def _run_blocks(self, h, t_emb, cond, rope_cos, rope_sin):
+        rope_cache = (rope_cos, rope_sin)
+        for block in self.blocks:
+            h = block(h, t_emb, cond, rope_cache=rope_cache)
+        return h
+
+    def _run_blocks_no_rope(self, h, t_emb, cond):
+        for block in self.blocks:
+            h = block(h, t_emb, cond, rope_cache=None)
+        return h
+
+    def __call__(
+        self,
+        x: MlxSparseTensor,
+        t: mx.array,
+        cond: mx.array,
+        concat_cond: MlxSparseTensor = None,
+    ) -> MlxSparseTensor:
+        if concat_cond is not None:
+            x = mlx_sparse_cat([x, concat_cond], dim=-1)
+
+        N = x.feats.shape[0]
+        logger.debug("[MLX] SLatFlow: N=%d, in_ch=%d, dtype=%s, mem=%s",
+                     N, x.feats.shape[1], x.feats.dtype, _metal_mem_mb())
+
+        h = self.input_layer(x.feats)
+
+        # Match upstream bfloat16 reduced-precision casting (manual_cast)
+        compute_dtype = mx.bfloat16
+        h = h.astype(compute_dtype)
+
+        t_emb = self.t_embedder(t)
+        if self.share_mod:
+            t_emb = self.adaLN_modulation(t_emb)
+        t_emb = t_emb.astype(compute_dtype)
+        cond = cond.astype(compute_dtype)
+
+        # Compute RoPE from sparse coords
+        if self.pe_mode == "rope":
+            coords_3d = x.coords[:, 1:]
+            rope_cos, rope_sin = compute_rope_phases(
+                coords_3d, self._rope_freqs,
+                self.model_channels // self.num_heads,
+            )
+
+            if self._compiled_blocks is None:
+                try:
+                    self._compiled_blocks = mx.compile(self._run_blocks)
+                    logger.info("[MLX] SLatFlow: using mx.compile for block loop")
+                except Exception as e:
+                    logger.warning("[MLX] SLatFlow: mx.compile failed (%s), falling back to per-block eval", e)
+                    self._compiled_blocks = False
+            if self._compiled_blocks:
+                h = self._compiled_blocks(h, t_emb, cond, rope_cos, rope_sin)
+            else:
+                rope_cache = (rope_cos, rope_sin)
+                for i, block in enumerate(self.blocks):
+                    h = block(h, t_emb, cond, rope_cache=rope_cache)
+                    if (i + 1) % 10 == 0:
+                        mx.eval(h)
+        else:
+            if self._compiled_blocks is None:
+                try:
+                    self._compiled_blocks = mx.compile(self._run_blocks_no_rope)
+                    logger.info("[MLX] SLatFlow: using mx.compile for block loop (no rope)")
+                except Exception as e:
+                    logger.warning("[MLX] SLatFlow: mx.compile failed (%s), falling back to per-block eval", e)
+                    self._compiled_blocks = False
+            if self._compiled_blocks:
+                h = self._compiled_blocks(h, t_emb, cond)
+            else:
+                for i, block in enumerate(self.blocks):
+                    h = block(h, t_emb, cond, rope_cache=None)
+                    if (i + 1) % 10 == 0:
+                        mx.eval(h)
+
+        logger.debug("[MLX] SLatFlow blocks done, N=%d, mem=%s", N, _metal_mem_mb())
+
+        # Cast back to float32 before final norm (matches upstream cast back to input dtype)
+        h = h.astype(mx.float32)
+        # Two-pass LayerNorm (matches LayerNorm32 / upstream F.layer_norm precision)
+        mean = mx.mean(h, axis=-1, keepdims=True)
+        h = h - mean
+        var = mx.mean(h * h, axis=-1, keepdims=True)
+        h = h * mx.rsqrt(var + 1e-5)
+        h = self.out_layer(h)
+        return x.replace(h)

mlx_backend/norm.py

@@ -0,0 +1,61 @@
+"""
+Normalization layers for MLX backend.
+Uses two-pass LayerNorm for parity with PyTorch (avoids fused kernel precision drift).
+"""
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class LayerNorm32(nn.Module):
+    """LayerNorm with two-pass variance for PyTorch parity.
+
+    mx.fast.layer_norm uses single-pass parallel variance that drifts ~1.4e-6
+    per call vs PyTorch's two-pass (~9.5e-7). Over 90+ norms per forward pass
+    and 50 sampler steps, this compounds significantly. Manual two-pass stays
+    in MLX's lazy graph while matching PyTorch numerics.
+    """
+
+    def __init__(self, dims: int, elementwise_affine: bool = True, eps: float = 1e-6):
+        super().__init__()
+        self.dims = dims
+        self.eps = eps
+        self.elementwise_affine = elementwise_affine
+        if elementwise_affine:
+            self.weight = mx.ones((dims,))
+            self.bias = mx.zeros((dims,))
+
+    def __call__(self, x: mx.array) -> mx.array:
+        x_dtype = x.dtype
+        x = x.astype(mx.float32)
+        mean = mx.mean(x, axis=-1, keepdims=True)
+        centered = x - mean
+        var = mx.mean(centered * centered, axis=-1, keepdims=True)
+        x = centered * mx.rsqrt(var + self.eps)
+        if self.elementwise_affine:
+            x = x * self.weight + self.bias
+        return x.astype(x_dtype)
+
+
+class SparseMultiHeadRMSNorm(nn.Module):
+    """Per-head RMSNorm using mx.fast.rms_norm (fused Metal kernel).
+
+    Equivalent to: L2_normalize(x) * gamma * sqrt(D)
+    Which equals: rms_norm(x) * gamma (the sqrt(D) cancels).
+    """
+
+    def __init__(self, dim: int, heads: int):
+        super().__init__()
+        self.dim = dim
+        self.scale = dim ** 0.5
+        self.gamma = mx.ones((heads, dim))
+
+    def __call__(self, x: mx.array) -> mx.array:
+        """x: (..., H, D) — supports (N, H, D) or (B, N, H, D)."""
+        x_dtype = x.dtype
+        orig_shape = x.shape
+        D = orig_shape[-1]
+        # Flatten all dims except last for fused rms_norm
+        x_flat = x.reshape(-1, D).astype(mx.float32)
+        x_flat = mx.fast.rms_norm(x_flat, None, 1e-6)
+        x = x_flat.reshape(orig_shape) * self.gamma
+        return x.astype(x_dtype)

mlx_backend/pipeline.py

@@ -0,0 +1,283 @@
+"""
+MLX pipeline factory — creates an upstream Trellis2ImageTo3DPipeline
+with MLX-backed model adapters.
+
+The upstream PT pipeline handles all orchestration, sampling, and mesh
+extraction. MLX models are injected via thin adapters that convert
+torch→mlx→torch at model boundaries.
+"""
+import os
+import gc
+import json
+import time
+import logging
+
+import mlx.core as mx
+
+logger = logging.getLogger(__name__)
+
+from . import load_safetensors, remap_flow_model_weights, remap_vae_decoder_weights
+from .flow_models import MlxSparseStructureFlowModel, MlxSLatFlowModel
+from .vae_decoders import MlxSparseUnetVaeDecoder, MlxFlexiDualGridVaeDecoder
+from .structure_decoder import load_structure_decoder
+from .dinov3 import load_dinov3_from_hf
+from .adapters import (
+    MlxFlowModelAdapter,
+    MlxStructureDecoderAdapter,
+    MlxFlexiDualGridAdapter,
+    MlxTexVaeDecoderAdapter,
+    MlxImageCondAdapter,
+)
+
+
+def _resolve_hf_path(rel_path: str) -> str:
+    """Resolve 'org/repo/path/to/file' to local HF cache path."""
+    from huggingface_hub import hf_hub_download
+    parts = rel_path.split('/')
+    repo_id = f"{parts[0]}/{parts[1]}"
+    file_base = '/'.join(parts[2:])
+    json_path = hf_hub_download(repo_id, f"{file_base}.json")
+    hf_hub_download(repo_id, f"{file_base}.safetensors")
+    return json_path.rsplit('.json', 1)[0]
+
+
+def _resolve_model_path(weights_path: str, rel_path: str) -> str:
+    """Resolve model path — local first, then HF Hub."""
+    full = os.path.join(weights_path, rel_path)
+    if os.path.exists(f"{full}.json"):
+        return full
+    return _resolve_hf_path(rel_path)
+
+
+def _load_mlx_flow_model(path: str, config: dict):
+    """Load an MLX flow model from config + safetensors."""
+    args = config['args']
+    if config['name'] == 'SparseStructureFlowModel':
+        model = MlxSparseStructureFlowModel(
+            resolution=args['resolution'],
+            in_channels=args['in_channels'],
+            model_channels=args['model_channels'],
+            cond_channels=args['cond_channels'],
+            out_channels=args['out_channels'],
+            num_blocks=args['num_blocks'],
+            num_heads=args.get('num_heads', 12),
+            mlp_ratio=args.get('mlp_ratio', 5.3334),
+            pe_mode=args.get('pe_mode', 'rope'),
+            share_mod=args.get('share_mod', True),
+            qk_rms_norm=args.get('qk_rms_norm', True),
+            qk_rms_norm_cross=args.get('qk_rms_norm_cross', True),
+        )
+        is_sparse = False
+    elif config['name'] in ('SLatFlowModel', 'ElasticSLatFlowModel'):
+        model = MlxSLatFlowModel(
+            resolution=args['resolution'],
+            in_channels=args['in_channels'],
+            model_channels=args['model_channels'],
+            cond_channels=args['cond_channels'],
+            out_channels=args['out_channels'],
+            num_blocks=args['num_blocks'],
+            num_heads=args.get('num_heads', 12),
+            mlp_ratio=args.get('mlp_ratio', 5.3334),
+            pe_mode=args.get('pe_mode', 'rope'),
+            share_mod=args.get('share_mod', True),
+            qk_rms_norm=args.get('qk_rms_norm', True),
+            qk_rms_norm_cross=args.get('qk_rms_norm_cross', True),
+        )
+        is_sparse = True
+    else:
+        raise ValueError(f"Unknown flow model type: {config['name']}")
+
+    weights = load_safetensors(f"{path}.safetensors")
+    weights = remap_flow_model_weights(weights)
+    model.load_weights(list(weights.items()))
+    return MlxFlowModelAdapter(model, is_sparse=is_sparse)
+
+
+def _load_mlx_structure_decoder(path: str, config: dict):
+    """Load MLX structure decoder and wrap in adapter."""
+    model = load_structure_decoder(path)
+    return MlxStructureDecoderAdapter(model)
+
+
+def _load_mlx_shape_decoder(path: str, config: dict):
+    """Load MLX FlexiDualGrid shape decoder and wrap in adapter."""
+    args = config['args']
+    model = MlxFlexiDualGridVaeDecoder(
+        resolution=args['resolution'],
+        model_channels=args['model_channels'],
+        latent_channels=args['latent_channels'],
+        num_blocks=args['num_blocks'],
+        block_type=args['block_type'],
+        up_block_type=args['up_block_type'],
+        block_args=args.get('block_args'),
+        use_fp16=args.get('use_fp16', False),
+    )
+    weights = load_safetensors(f"{path}.safetensors")
+    weights = remap_vae_decoder_weights(weights)
+    weights = {f'decoder.{k}': v for k, v in weights.items()}
+    model.load_weights(list(weights.items()))
+    return MlxFlexiDualGridAdapter(model)
+
+
+def _load_mlx_tex_decoder(path: str, config: dict):
+    """Load MLX texture VAE decoder and wrap in adapter."""
+    args = config['args']
+    model = MlxSparseUnetVaeDecoder(
+        out_channels=args['out_channels'],
+        model_channels=args['model_channels'],
+        latent_channels=args['latent_channels'],
+        num_blocks=args['num_blocks'],
+        block_type=args['block_type'],
+        up_block_type=args['up_block_type'],
+        block_args=args.get('block_args'),
+        use_fp16=args.get('use_fp16', False),
+        pred_subdiv=args.get('pred_subdiv', True),
+    )
+    weights = load_safetensors(f"{path}.safetensors")
+    weights = remap_vae_decoder_weights(weights)
+    model.load_weights(list(weights.items()))
+    return MlxTexVaeDecoderAdapter(model)
+
+
+# Map model name patterns to loader functions
+_LOADER_MAP = {
+    'sparse_structure_flow_model': _load_mlx_flow_model,
+    'sparse_structure_decoder': _load_mlx_structure_decoder,
+    'shape_slat_decoder': _load_mlx_shape_decoder,
+    'tex_slat_decoder': _load_mlx_tex_decoder,
+}
+
+
+def _get_loader(name: str, config: dict):
+    """Pick the right loader for a model name."""
+    # Exact match first
+    if name in _LOADER_MAP:
+        return _LOADER_MAP[name]
+    # Flow models by name pattern
+    if 'flow_model' in name:
+        return _load_mlx_flow_model
+    # VAE decoders by config name
+    if config['name'] == 'FlexiDualGridVaeDecoder':
+        return _load_mlx_shape_decoder
+    if config['name'] == 'SparseUnetVaeDecoder':
+        return _load_mlx_tex_decoder
+    raise ValueError(f"No loader for model '{name}' (type: {config['name']})")
+
+
+def create_mlx_pipeline(weights_path: str = "weights/TRELLIS.2-4B"):
+    """Create upstream Trellis2ImageTo3DPipeline with MLX-backed models.
+
+    All model compute runs in MLX. The upstream PT pipeline handles
+    orchestration, sampling (FlowEulerCfgSampler etc.), and mesh extraction.
+    """
+    import torch
+    from trellis2.pipelines.trellis2_image_to_3d import Trellis2ImageTo3DPipeline
+    from trellis2.pipelines import samplers
+    from trellis2.pipelines.rembg import BiRefNet
+
+    print(f"[MLX] Loading pipeline config from {weights_path}...")
+    config_file = os.path.join(weights_path, "pipeline.json")
+    with open(config_file) as f:
+        args = json.load(f)['args']
+
+    # Load all models with MLX adapters
+    models = {}
+    for name, rel_path in args['models'].items():
+        path = _resolve_model_path(weights_path, rel_path)
+        with open(f"{path}.json") as f:
+            model_config = json.load(f)
+
+        t0 = time.time()
+        loader = _get_loader(name, model_config)
+        models[name] = loader(path, model_config)
+        dt = time.time() - t0
+        print(f"  [MLX] Loaded '{name}' in {dt:.1f}s")
+
+    # Create upstream pipeline with MLX models
+    pipeline = Trellis2ImageTo3DPipeline(models)
+    pipeline._pretrained_args = args
+
+    # Set up samplers (upstream PT classes — source of truth for sampling)
+    pipeline.sparse_structure_sampler = getattr(
+        samplers, args['sparse_structure_sampler']['name']
+    )(**args['sparse_structure_sampler']['args'])
+    pipeline.sparse_structure_sampler_params = args['sparse_structure_sampler']['params']
+
+    pipeline.shape_slat_sampler = getattr(
+        samplers, args['shape_slat_sampler']['name']
+    )(**args['shape_slat_sampler']['args'])
+    pipeline.shape_slat_sampler_params = args['shape_slat_sampler']['params']
+
+    pipeline.tex_slat_sampler = getattr(
+        samplers, args['tex_slat_sampler']['name']
+    )(**args['tex_slat_sampler']['args'])
+    pipeline.tex_slat_sampler_params = args['tex_slat_sampler']['params']
+
+    # Normalization
+    pipeline.shape_slat_normalization = args['shape_slat_normalization']
+    pipeline.tex_slat_normalization = args['tex_slat_normalization']
+
+    # Image conditioning (MLX DINOv3)
+    pipeline.image_cond_model = MlxImageCondAdapter(
+        load_dinov3_from_hf(args['image_cond_model']['args']['model_name'])
+    )
+
+    # Background removal (PT — lightweight, used once)
+    pipeline.rembg_model = BiRefNet(**args['rembg_model']['args'])
+
+    pipeline.low_vram = True
+    pipeline._device = torch.device('cpu')
+    pipeline.default_pipeline_type = args.get('default_pipeline_type', '1024_cascade')
+    pipeline.pbr_attr_layout = {
+        'base_color': slice(0, 3),
+        'metallic': slice(3, 4),
+        'roughness': slice(4, 5),
+        'alpha': slice(5, 6),
+    }
+
+    print("[MLX] Pipeline ready.")
+    return pipeline
+
+
+def to_glb(mesh, output_path: str,
+           decimation_target: int = 1000000,
+           texture_size: int = 2048,
+           remesh: bool = False,
+           verbose: bool = True) -> str:
+    """Export MeshWithVoxel to GLB file."""
+    import o_voxel
+
+    print(f"Exporting to {output_path}...")
+    glb = o_voxel.postprocess.to_glb(
+        vertices=mesh.vertices,
+        faces=mesh.faces,
+        attr_volume=mesh.attrs,
+        coords=mesh.coords,
+        attr_layout=mesh.layout,
+        voxel_size=mesh.voxel_size,
+        aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
+        decimation_target=decimation_target,
+        texture_size=texture_size,
+        remesh=remesh,
+        verbose=verbose,
+    )
+    glb.export(output_path)
+    print(f"Exported: {output_path}")
+    return output_path
+
+
+# Backward-compat alias
+class MlxTrellis2Pipeline:
+    """Deprecated — use create_mlx_pipeline() instead.
+
+    Thin wrapper that creates the upstream pipeline and delegates .run()/.to_glb().
+    """
+
+    def __init__(self, weights_path: str = "weights/TRELLIS.2-4B"):
+        self._pipeline = create_mlx_pipeline(weights_path)
+
+    def run(self, image, **kwargs):
+        return self._pipeline.run(image, **kwargs)
+
+    def to_glb(self, mesh, output_path, **kwargs):
+        return to_glb(mesh, output_path, **kwargs)

mlx_backend/rope.py

@@ -0,0 +1,96 @@
+"""
+3D Rotary Position Embeddings for MLX.
+Uses a custom Metal kernel for fused application.
+"""
+import mlx.core as mx
+import math
+
+
+def build_rope_freqs(head_dim: int, dim: int = 3,
+                     rope_freq: tuple = (1.0, 10000.0)) -> mx.array:
+    """
+    Build frequency table for RoPE.
+
+    Returns:
+        freqs: (freq_dim,) array of frequencies
+    """
+    freq_dim = head_dim // 2 // dim
+    freqs = mx.arange(freq_dim, dtype=mx.float32) / freq_dim
+    freqs = rope_freq[0] / (rope_freq[1] ** freqs)
+    return freqs
+
+
+def compute_rope_phases(coords: mx.array, freqs: mx.array, head_dim: int) -> mx.array:
+    """
+    Compute RoPE cos/sin phases from 3D coordinates.
+
+    Args:
+        coords: (N, 3) int coordinates
+        freqs: (freq_dim,) frequency table
+        head_dim: head dimension
+
+    Returns:
+        cos_phases: (N, head_dim//2) cos values
+        sin_phases: (N, head_dim//2) sin values
+    """
+    N = coords.shape[0]
+    dim = coords.shape[1]
+    freq_dim = freqs.shape[0]
+    target = head_dim // 2
+
+    # Vectorized: compute all phases in one matmul-like op
+    # coords: (N, 3) float, freqs: (freq_dim,)
+    # phases[d] = coords[:, d:d+1] * freqs -> (N, freq_dim) per dim
+    coords_f = coords.astype(mx.float32)
+    phases_list = []
+    for d in range(dim):
+        phases_list.append(coords_f[:, d:d+1] * freqs[None, :])
+    phases = mx.concatenate(phases_list, axis=-1)  # (N, dim * freq_dim)
+
+    if phases.shape[-1] < target:
+        pad_n = target - phases.shape[-1]
+        phases = mx.concatenate([phases, mx.zeros((N, pad_n))], axis=-1)
+
+    cos_phases = mx.cos(phases)
+    sin_phases = mx.sin(phases)
+    return cos_phases, sin_phases
+
+
+def apply_rope(x: mx.array, cos: mx.array, sin: mx.array) -> mx.array:
+    """
+    Apply rotary embeddings using interleaved pairing (matching PyTorch complex multiply).
+
+    PyTorch pairs adjacent elements: (x[0],x[1]), (x[2],x[3]), ...
+    via view_as_complex → complex multiply → view_as_real.
+
+    Args:
+        x: (..., N, H, D) features — (N, H, D) or (B, N, H, D)
+        cos: (N, D//2) cos phases
+        sin: (N, D//2) sin phases
+
+    Returns:
+        Same shape as x, rotated.
+    """
+    # Interleaved: pair (x[0],x[1]), (x[2],x[3]), etc.
+    # Reshape last dim from D to (D//2, 2)
+    orig_shape = x.shape
+    x_paired = x.reshape(*orig_shape[:-1], -1, 2)  # (..., D//2, 2)
+    x1 = x_paired[..., 0]  # even indices: (..., D//2)
+    x2 = x_paired[..., 1]  # odd indices:  (..., D//2)
+
+    if x.ndim == 3:
+        # (N, H, D//2) — expand cos/sin to (N, 1, D//2)
+        cos = cos[:, None, :]
+        sin = sin[:, None, :]
+    else:
+        # (B, N, H, D//2) — expand cos/sin to (1, N, 1, D//2)
+        cos = cos[None, :, None, :]
+        sin = sin[None, :, None, :]
+
+    # Complex multiply: (x1 + i*x2) * (cos + i*sin)
+    o1 = x1 * cos - x2 * sin  # real part
+    o2 = x1 * sin + x2 * cos  # imag part
+
+    # Interleave back: stack on last dim then flatten
+    out = mx.stack([o1, o2], axis=-1)  # (..., D//2, 2)
+    return out.reshape(orig_shape)

mlx_backend/sparse_conv.py

@@ -0,0 +1,170 @@
+"""
+Submanifold sparse 3D convolution for MLX.
+Port of conv_pytorch.py algorithm: hash → neighbor map → gather → bmm → scatter.
+"""
+import itertools
+import logging
+import time
+import mlx.core as mx
+import mlx.nn as nn
+from .sparse_tensor import MlxSparseTensor
+
+logger = logging.getLogger(__name__)
+
+
+def _build_neighbor_map(
+    coords: mx.array,
+    batch_size: int,
+    spatial_shape: tuple,
+    kernel_size: tuple,
+    dilation: tuple,
+) -> mx.array:
+    """
+    Build neighbor map for submanifold sparse conv.
+
+    Returns:
+        neighbor_map: (K, N) int32. neighbor_map[k, i] = index of neighbor, or N (pad index).
+    """
+    N = coords.shape[0]
+    D, H, W = spatial_shape
+    DHW = D * H * W
+
+    # Build lookup table: flat_coord -> index
+    flat_keys = (coords[:, 0].astype(mx.int32) * DHW +
+                 coords[:, 1].astype(mx.int32) * (H * W) +
+                 coords[:, 2].astype(mx.int32) * W +
+                 coords[:, 3].astype(mx.int32))
+
+    table_size = batch_size * DHW
+    # Initialize lookup with N (= pad index)
+    lookup = mx.full((table_size,), N, dtype=mx.int32)
+    lookup = lookup.at[flat_keys].add(mx.arange(N, dtype=mx.int32) - N)
+
+    # Generate kernel offsets
+    kd, kh, kw = kernel_size
+    dd, dh, dw = dilation
+    offsets = []
+    for dx, dy, dz in itertools.product(
+        range(-(kd // 2), kd // 2 + 1),
+        range(-(kh // 2), kh // 2 + 1),
+        range(-(kw // 2), kw // 2 + 1),
+    ):
+        offsets.append((dx * dd, dy * dh, dz * dw))
+    K = len(offsets)
+
+    # For each offset, shift coords and lookup
+    neighbor_maps = []
+    for dx, dy, dz in offsets:
+        sx = coords[:, 1].astype(mx.int32) + dx
+        sy = coords[:, 2].astype(mx.int32) + dy
+        sz = coords[:, 3].astype(mx.int32) + dz
+
+        # Bounds check
+        valid = ((sx >= 0) & (sx < D) &
+                 (sy >= 0) & (sy < H) &
+                 (sz >= 0) & (sz < W))
+
+        flat_shifted = (coords[:, 0].astype(mx.int32) * DHW +
+                        sx * (H * W) + sy * W + sz)
+        # Clamp for safe indexing
+        flat_shifted = mx.clip(flat_shifted, 0, table_size - 1)
+        looked_up = lookup[flat_shifted]
+        # Set invalid to N (pad index)
+        looked_up = mx.where(valid, looked_up, N)
+        neighbor_maps.append(looked_up)
+
+    return mx.stack(neighbor_maps, axis=0)  # (K, N)
+
+
+class MlxSparseConv3d(nn.Module):
+    """
+    Submanifold sparse 3D convolution in MLX.
+
+    Weight format: (Co, Kd, Kh, Kw, Ci) — same as PyTorch checkpoint.
+    At forward time, reshaped to (K, Ci, Co) for gather-matmul pattern.
+    """
+
+    def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3,
+                 dilation: int = 1):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        if isinstance(kernel_size, int):
+            self.kernel_size = (kernel_size,) * 3
+        else:
+            self.kernel_size = tuple(kernel_size)
+        if isinstance(dilation, int):
+            self.dilation = (dilation,) * 3
+        else:
+            self.dilation = tuple(dilation)
+
+        K = self.kernel_size[0] * self.kernel_size[1] * self.kernel_size[2]
+        # Weight stored as (Co, Kd, Kh, Kw, Ci) for checkpoint compatibility
+        self.weight = mx.zeros((out_channels, *self.kernel_size, in_channels))
+        self.bias = mx.zeros((out_channels,))
+
+    def __call__(self, x: MlxSparseTensor) -> MlxSparseTensor:
+        t0 = time.time()
+        N = x.feats.shape[0]
+        Co = self.out_channels
+        Ci = self.in_channels
+
+        # Get or build neighbor map
+        cache_key = f'SubMConv3d_neighbor_cache_mlx_{self.kernel_size}_dilation{self.dilation}'
+        neighbor_map = x.get_spatial_cache(cache_key)
+
+        if neighbor_map is None:
+            batch_size = x.shape[0]
+            spatial_shape = x.spatial_shape
+            neighbor_map = _build_neighbor_map(
+                x.coords, batch_size, spatial_shape,
+                self.kernel_size, self.dilation,
+            )
+            mx.eval(neighbor_map)  # eval now to free the O(D³) lookup table
+            x.register_spatial_cache(cache_key, neighbor_map)
+
+        K = neighbor_map.shape[0]
+
+        # Reshape weight: (Co, Kd, Kh, Kw, Ci) -> (K, Ci, Co)
+        w = self.weight.reshape(Co, K, Ci)  # (Co, K, Ci)
+        w = w.transpose(1, 2, 0)  # (K, Ci, Co)
+
+        # Pad feats with zero row for out-of-bounds indices
+        # feats_padded[N] is zeros, so invalid neighbor lookups produce zero
+        # from matmul naturally — no valid_mask needed.
+        feats_padded = mx.concatenate([
+            x.feats,
+            mx.zeros((1, Ci), dtype=x.feats.dtype)
+        ], axis=0)  # (N+1, Ci)
+
+        # Target 512MB chunks — MLX handles chunked graphs well and prevents OOM
+        max_bytes = 512 * 1024**2
+        elem_size = 2 if w.dtype in (mx.float16, mx.bfloat16) else 4
+        per_offset_bytes = N * max(Ci, Co) * elem_size * 2  # gather + matmul output
+        chunk_size = max(1, min(K, int(max_bytes / max(per_offset_bytes, 1))))
+
+        logger.debug("[MLX] SparseConv3d: N=%d, Ci=%d, Co=%d, K=%d, chunk_size=%d",
+                     N, Ci, Co, K, chunk_size)
+
+        result = mx.zeros((N, Co), dtype=w.dtype)
+        if chunk_size >= K:
+            # Single pass — no eval needed, keep graph lazy for caller
+            gathered = feats_padded[neighbor_map]  # (K, N, Ci)
+            out = mx.matmul(gathered.astype(w.dtype), w)  # (K, N, Co)
+            result = mx.sum(out, axis=0)  # (N, Co)
+        else:
+            for start in range(0, K, chunk_size):
+                end = min(start + chunk_size, K)
+                nmap_chunk = neighbor_map[start:end]  # (chunk, N)
+                gathered = feats_padded[nmap_chunk]  # (chunk, N, Ci)
+                out = mx.matmul(gathered.astype(w.dtype), w[start:end])  # (chunk, N, Co)
+                result = result + mx.sum(out, axis=0)  # (N, Co)
+                mx.eval(result)  # free chunk intermediates between iterations
+
+        result = result + self.bias
+
+        dt = time.time() - t0
+        if dt > 0.1:
+            logger.debug("[MLX] SparseConv3d done: N=%d, %dx%d→%d, %.2fs", N, Ci, K, Co, dt)
+
+        return x.replace(result.astype(x.feats.dtype))

mlx_backend/sparse_ops.py

@@ -0,0 +1,284 @@
+"""
+Sparse spatial operations for MLX: Downsample, Upsample, Channel2Spatial, Spatial2Channel.
+"""
+import logging
+import time
+import numpy as np
+import mlx.core as mx
+import mlx.nn as nn
+from .sparse_tensor import MlxSparseTensor
+
+logger = logging.getLogger(__name__)
+
+
+class MlxSparseDownsample(nn.Module):
+    """Downsample sparse tensor by factor, using mean pooling."""
+
+    def __init__(self, factor: int = 2):
+        super().__init__()
+        self.factor = factor
+
+    def __call__(self, x: MlxSparseTensor) -> MlxSparseTensor:
+        t0 = time.time()
+        N_in = x.feats.shape[0]
+
+        cache = x.get_spatial_cache(f'downsample_{self.factor}')
+        if cache is None:
+            DIM = x.coords.shape[-1] - 1  # 3
+            coords = x.coords  # (N, 4)
+            batch_col = coords[:, 0]
+            spatial_cols = [coords[:, i + 1] // self.factor for i in range(DIM)]
+
+            MAX = [(int(x.spatial_shape[i]) + self.factor - 1) // self.factor for i in range(DIM)]
+            OFFSET = [1] * (DIM + 1)
+            for i in range(DIM - 1, -1, -1):
+                OFFSET[i] = OFFSET[i + 1] * MAX[i]
+            # batch offset
+            batch_offset = OFFSET[0]
+
+            code = batch_col.astype(mx.int32) * batch_offset
+            for i in range(DIM):
+                code = code + spatial_cols[i].astype(mx.int32) * OFFSET[i + 1]
+
+            unique_codes, idx = mx.unique(code, return_inverse=True)
+
+            # Reconstruct new coords from unique codes
+            new_batch = unique_codes // batch_offset
+            remainder = unique_codes % batch_offset
+            new_spatial = []
+            for i in range(DIM):
+                new_spatial.append(remainder // OFFSET[i + 1])
+                remainder = remainder % OFFSET[i + 1]
+            new_coords = mx.stack([new_batch] + new_spatial, axis=-1).astype(mx.int32)
+        else:
+            new_coords, idx = cache
+
+        # Scatter-mean: vectorized scatter-add + counts
+        N_new = new_coords.shape[0]
+        C = x.feats.shape[1]
+        new_feats = _mlx_scatter_mean(x.feats, idx, N_new)
+
+        out = MlxSparseTensor(new_feats, new_coords, shape=(x.shape[0], C))
+        out._scale = tuple(s * self.factor for s in x._scale)
+        out._spatial_cache = x._spatial_cache
+
+        if cache is None:
+            x.register_spatial_cache(f'downsample_{self.factor}', (new_coords, idx))
+            out.register_spatial_cache(f'upsample_{self.factor}', (x.coords, idx))
+            out.register_spatial_cache('shape', tuple(MAX))
+
+        dt = time.time() - t0
+        logger.debug("[MLX] Downsample: N %d → %d (factor=%d), %.2fs", N_in, N_new, self.factor, dt)
+        return out
+
+
+def _mlx_scatter_mean(feats: mx.array, idx: mx.array, n_out: int) -> mx.array:
+    """
+    Scatter-mean using vectorized scatter-add.
+    feats: (N, C), idx: (N,) -> (n_out, C)
+    """
+    C = feats.shape[1]
+    feats_f32 = feats.astype(mx.float32)
+
+    # Scatter-add features and counts using .at[].add()
+    sums = mx.zeros((n_out, C), dtype=mx.float32)
+    counts = mx.zeros((n_out, 1), dtype=mx.float32)
+    idx_int = idx.astype(mx.int32)
+    sums = sums.at[idx_int].add(feats_f32)
+    counts = counts.at[idx_int].add(mx.ones((feats.shape[0], 1), dtype=mx.float32))
+    return (sums / mx.maximum(counts, 1.0)).astype(feats.dtype)
+
+
+class MlxSparseUpsample(nn.Module):
+    """Upsample sparse tensor by factor using nearest neighbor."""
+
+    def __init__(self, factor: int = 2):
+        super().__init__()
+        self.factor = factor
+
+    def __call__(self, x: MlxSparseTensor, subdivision: MlxSparseTensor = None) -> MlxSparseTensor:
+        t0 = time.time()
+        N_in = x.feats.shape[0]
+        DIM = x.coords.shape[-1] - 1  # 3
+        cache = x.get_spatial_cache(f'upsample_{self.factor}')
+
+        if cache is None:
+            if subdivision is None:
+                raise ValueError('Cache not found. Provide subdivision tensor.')
+            sub = subdivision.feats  # (N, factor^DIM)
+            F_DIM = self.factor ** DIM
+
+            # Vectorized: find active sub-voxels via numpy nonzero
+            sub_flat = sub.reshape(-1)
+            mx.eval(sub_flat)
+            active_np = np.where(np.array(sub_flat) > 0)[0].astype(np.int32)
+            active_indices = mx.array(active_np)
+
+            parent_idx = active_indices // F_DIM
+            subidx = active_indices % F_DIM
+
+            parent_coords = x.coords[parent_idx].astype(mx.int32)
+            batch_col = parent_coords[:, :1]
+            spatial = parent_coords[:, 1:] * self.factor
+            for d in range(DIM):
+                offset = (subidx // (self.factor ** d) % self.factor).astype(mx.int32)[:, None]
+                spatial = mx.concatenate([
+                    spatial[:, :d],
+                    spatial[:, d:d+1] + offset,
+                    spatial[:, d+1:],
+                ], axis=-1)
+            new_coords = mx.concatenate([batch_col, spatial], axis=-1)
+            idx = parent_idx
+        else:
+            new_coords, idx = cache
+
+        new_feats = x.feats[idx]
+        out = MlxSparseTensor(new_feats, new_coords, shape=(x.shape[0], x.feats.shape[1]))
+        out._scale = tuple(s / self.factor for s in x._scale)
+        if cache is not None:
+            out._spatial_cache = x._spatial_cache
+
+        dt = time.time() - t0
+        N_out = new_feats.shape[0]
+        logger.debug("[MLX] Upsample: N %d → %d (factor=%d), %.2fs", N_in, N_out, self.factor, dt)
+        return out
+
+
+class MlxSparseSpatial2Channel(nn.Module):
+    """Downsample by rearranging spatial dims into channels."""
+
+    def __init__(self, factor: int = 2):
+        super().__init__()
+        self.factor = factor
+
+    def __call__(self, x: MlxSparseTensor) -> MlxSparseTensor:
+        t0 = time.time()
+        N_in = x.feats.shape[0]
+        DIM = x.coords.shape[-1] - 1
+        F = self.factor
+        cache = x.get_spatial_cache(f'spatial2channel_{F}')
+
+        if cache is None:
+            coords = x.coords
+            batch_col = coords[:, 0]
+            spatial_cols = [coords[:, i + 1] // F for i in range(DIM)]
+
+            # Sub-index within the factor cube
+            subidx_parts = [coords[:, i + 1] % F for i in range(DIM)]
+            subidx = subidx_parts[0].astype(mx.int32)
+            for d in range(1, DIM):
+                subidx = subidx + subidx_parts[d].astype(mx.int32) * (F ** d)
+
+            MAX = [(int(x.spatial_shape[i]) + F - 1) // F for i in range(DIM)]
+            OFFSET = [1] * (DIM + 1)
+            for i in range(DIM - 1, -1, -1):
+                OFFSET[i] = OFFSET[i + 1] * MAX[i]
+            batch_offset = OFFSET[0]
+
+            code = batch_col.astype(mx.int32) * batch_offset
+            for i in range(DIM):
+                code = code + spatial_cols[i].astype(mx.int32) * OFFSET[i + 1]
+
+            unique_codes, idx = mx.unique(code, return_inverse=True)
+
+            new_batch = unique_codes // batch_offset
+            remainder = unique_codes % batch_offset
+            new_spatial = []
+            for i in range(DIM):
+                new_spatial.append(remainder // OFFSET[i + 1])
+                remainder = remainder % OFFSET[i + 1]
+            new_coords = mx.stack([new_batch] + new_spatial, axis=-1).astype(mx.int32)
+        else:
+            new_coords, idx, subidx = cache
+
+        # Pack features: scatter into (N_new * F^DIM, C) then reshape to (N_new, C * F^DIM)
+        N_new = new_coords.shape[0]
+        C = x.feats.shape[1]
+        F_DIM = F ** DIM
+        new_feats = mx.zeros((N_new * F_DIM, C), dtype=x.feats.dtype)
+        flat_idx = idx * F_DIM + subidx.astype(mx.int32)
+        new_feats = new_feats.at[flat_idx].add(x.feats)
+        new_feats = new_feats.reshape(N_new, C * F_DIM)
+
+        out = MlxSparseTensor(
+            new_feats, new_coords,
+            shape=(x.shape[0], C * F_DIM) if x.shape is not None else None,
+        )
+        out._scale = tuple(s * F for s in x._scale)
+        out._spatial_cache = x._spatial_cache
+
+        if cache is None:
+            x.register_spatial_cache(f'spatial2channel_{F}', (new_coords, idx, subidx))
+            out.register_spatial_cache(f'channel2spatial_{F}', (x.coords, idx, subidx))
+            out.register_spatial_cache('shape', tuple(MAX))
+
+        dt = time.time() - t0
+        logger.debug("[MLX] Spatial2Channel: N %d → %d, C %d → %d, %.2fs",
+                     N_in, N_new, C, C * F_DIM, dt)
+        return out
+
+
+class MlxSparseChannel2Spatial(nn.Module):
+    """Upsample by rearranging channels into spatial dims."""
+
+    def __init__(self, factor: int = 2):
+        super().__init__()
+        self.factor = factor
+
+    def __call__(self, x: MlxSparseTensor, subdivision: MlxSparseTensor = None) -> MlxSparseTensor:
+        t0 = time.time()
+        N_in = x.feats.shape[0]
+        DIM = x.coords.shape[-1] - 1
+        F = self.factor
+        F_DIM = F ** DIM
+        cache = x.get_spatial_cache(f'channel2spatial_{F}')
+
+        if cache is None:
+            if subdivision is None:
+                raise ValueError('Cache not found. Provide subdivision tensor.')
+            sub = subdivision.feats  # (N, F_DIM)
+
+            # Vectorized: find active sub-voxels via numpy nonzero
+            sub_flat = sub.reshape(-1)
+            mx.eval(sub_flat)
+            active_np = np.where(np.array(sub_flat) > 0)[0].astype(np.int32)
+            active_indices = mx.array(active_np)
+
+            parent_idx = active_indices // F_DIM
+            subidx = active_indices % F_DIM
+
+            parent_coords = x.coords[parent_idx].astype(mx.int32)
+            batch_col = parent_coords[:, :1]
+            spatial = parent_coords[:, 1:] * F
+            for d in range(DIM):
+                offset = (subidx // (F ** d) % F).astype(mx.int32)[:, None]
+                spatial = mx.concatenate([
+                    spatial[:, :d],
+                    spatial[:, d:d+1] + offset,
+                    spatial[:, d+1:],
+                ], axis=-1)
+            new_coords = mx.concatenate([batch_col, spatial], axis=-1)
+            idx = parent_idx
+        else:
+            new_coords, idx, subidx = cache
+
+        # Unpack: reshape (N, C * F^DIM) -> (N * F^DIM, C), then gather
+        C_packed = x.feats.shape[1]
+        C = C_packed // F_DIM
+        x_feats = x.feats.reshape(x.feats.shape[0] * F_DIM, C)
+        flat_idx = idx * F_DIM + subidx.astype(mx.int32)
+        new_feats = x_feats[flat_idx]
+
+        out = MlxSparseTensor(
+            new_feats, new_coords,
+            shape=(x.shape[0], C) if x.shape is not None else None,
+        )
+        out._scale = tuple(s / F for s in x._scale)
+        if cache is not None:
+            out._spatial_cache = x._spatial_cache
+
+        dt = time.time() - t0
+        N_out = new_feats.shape[0]
+        logger.debug("[MLX] Channel2Spatial: N %d → %d, C %d → %d, %.2fs",
+                     N_in, N_out, C_packed, C, dt)
+        return out

mlx_backend/sparse_tensor.py

@@ -0,0 +1,124 @@
+"""
+MLX sparse tensor representation.
+Mirrors trellis2.modules.sparse.basic.SparseTensor but uses mx.array.
+"""
+from dataclasses import dataclass, field
+from typing import Optional, Dict, Any, List, Tuple
+from fractions import Fraction
+import mlx.core as mx
+
+
+@dataclass
+class MlxSparseTensor:
+    """
+    Sparse tensor with MLX arrays.
+
+    Fields:
+        feats: (N, C) feature matrix
+        coords: (N, 4) int32 coordinates [batch, x, y, z]
+        shape: (B, C) batch shape
+        _scale: coordinate scale factors (for cache keying)
+        _spatial_cache: nested dict keyed by scale, then by cache name
+    """
+    feats: mx.array
+    coords: mx.array
+    shape: Optional[Tuple[int, ...]] = None
+    _scale: Tuple = (Fraction(1, 1), Fraction(1, 1), Fraction(1, 1))
+    _spatial_cache: Dict[str, Dict[str, Any]] = field(default_factory=dict)
+
+    def __post_init__(self):
+        if self.shape is None:
+            batch_max = int(self.coords[:, 0].max().item()) + 1
+            self.shape = (batch_max, *self.feats.shape[1:])
+
+    @property
+    def device(self):
+        return self.feats.dtype  # MLX doesn't have device, everything is on GPU
+
+    @property
+    def dtype(self):
+        return self.feats.dtype
+
+    @property
+    def N(self) -> int:
+        return self.feats.shape[0]
+
+    @property
+    def spatial_shape(self) -> Tuple[int, ...]:
+        cached = self.get_spatial_cache('shape')
+        if cached is not None:
+            return cached
+        maxes = self.coords[:, 1:].max(axis=0)
+        ss = tuple((int(m) + 1) for m in maxes.tolist())
+        self.register_spatial_cache('shape', ss)
+        return ss
+
+    @property
+    def layout(self) -> List[slice]:
+        cached = self.get_spatial_cache('layout')
+        if cached is not None:
+            return cached
+        batch_size = self.shape[0]
+        batch_ids = self.coords[:, 0]
+        layout = []
+        start = 0
+        for b in range(batch_size):
+            count = int(mx.sum(batch_ids == b).item())
+            layout.append(slice(start, start + count))
+            start += count
+        self.register_spatial_cache('layout', layout)
+        return layout
+
+    def replace(self, feats: mx.array, coords: Optional[mx.array] = None) -> 'MlxSparseTensor':
+        """Create a new MlxSparseTensor with new features (and optionally new coords)."""
+        new_shape = (self.shape[0], *feats.shape[1:]) if self.shape is not None else None
+        return MlxSparseTensor(
+            feats=feats,
+            coords=coords if coords is not None else self.coords,
+            shape=new_shape,
+            _scale=self._scale,
+            _spatial_cache=self._spatial_cache,
+        )
+
+    def register_spatial_cache(self, key: str, value: Any) -> None:
+        scale_key = str(self._scale)
+        if scale_key not in self._spatial_cache:
+            self._spatial_cache[scale_key] = {}
+        self._spatial_cache[scale_key][key] = value
+
+    def get_spatial_cache(self, key: str = None) -> Any:
+        scale_key = str(self._scale)
+        cur = self._spatial_cache.get(scale_key, {})
+        if key is None:
+            return cur
+        return cur.get(key, None)
+
+    def __repr__(self):
+        return f"MlxSparseTensor(N={self.N}, shape={self.shape}, dtype={self.dtype})"
+
+
+def mlx_sparse_cat(inputs: List[MlxSparseTensor], dim: int = -1) -> MlxSparseTensor:
+    """Concatenate sparse tensors along feature dimension."""
+    if dim == -1 or dim == 1:
+        feats = mx.concatenate([s.feats for s in inputs], axis=-1)
+        return inputs[0].replace(feats)
+    elif dim == 0:
+        # Batch concatenation: shift batch indices
+        offset = 0
+        all_coords = []
+        all_feats = []
+        for s in inputs:
+            coords = mx.array(s.coords)  # copy
+            coords = mx.concatenate([
+                coords[:, :1] + offset,
+                coords[:, 1:]
+            ], axis=-1)
+            all_coords.append(coords)
+            all_feats.append(s.feats)
+            offset += s.shape[0]
+        return MlxSparseTensor(
+            feats=mx.concatenate(all_feats, axis=0),
+            coords=mx.concatenate(all_coords, axis=0),
+        )
+    else:
+        raise ValueError(f"Unsupported dim={dim}")

mlx_backend/structure_decoder.py

@@ -0,0 +1,214 @@
+"""
+Sparse structure decoder in MLX — dense 3D conv network.
+Converts flow output (1, 8, 16, 16, 16) → binary voxel grid (1, 1, 64, 64, 64).
+"""
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+
+
+def conv3d(x: mx.array, weight: mx.array, bias: mx.array = None, padding: int = 1) -> mx.array:
+    """3D convolution via mx.conv_general. x: (B,D,H,W,Ci), weight: (Co,Ci,kD,kH,kW)."""
+    # MLX conv_general expects (B, spatial..., Ci) and weight (Co, kD, kH, kW, Ci)
+    # PyTorch weight format: (Co, Ci, kD, kH, kW) → MLX: (Co, kD, kH, kW, Ci)
+    w = weight.transpose(0, 2, 3, 4, 1)  # (Co, kD, kH, kW, Ci)
+    out = mx.conv_general(x, w, padding=padding)
+    if bias is not None:
+        out = out + bias
+    return out
+
+
+class ChannelLayerNorm3d(nn.Module):
+    """
+    Channel LayerNorm for 3D data in channels-last format.
+    Matches PT ChannelLayerNorm32: LayerNorm applied to the channel dimension.
+    Input: (B, D, H, W, C) — normalizes over C.
+    """
+
+    def __init__(self, channels: int):
+        super().__init__()
+        self.weight = mx.ones((channels,))
+        self.bias = mx.zeros((channels,))
+
+    def __call__(self, x: mx.array) -> mx.array:
+        x_dtype = x.dtype
+        x = x.astype(mx.float32)
+        mean = mx.mean(x, axis=-1, keepdims=True)
+        centered = x - mean
+        var = mx.mean(centered * centered, axis=-1, keepdims=True)
+        x = centered * mx.rsqrt(var + 1e-5) * self.weight + self.bias
+        return x.astype(x_dtype)
+
+
+class ResBlock3d(nn.Module):
+    """ResBlock: ChannelLayerNorm → SiLU → Conv3d → ChannelLayerNorm → SiLU → Conv3d + skip.
+
+    Uses ChannelLayerNorm to match PT's default norm_type="layer" (ChannelLayerNorm32).
+    """
+
+    def __init__(self, channels: int):
+        super().__init__()
+        self.channels = channels
+        self.norm1 = ChannelLayerNorm3d(channels)
+        self.norm2 = ChannelLayerNorm3d(channels)
+        # Conv weights stored in PyTorch format (Co, Ci, kD, kH, kW)
+        self.conv1 = {"weight": mx.zeros((channels, channels, 3, 3, 3)),
+                      "bias": mx.zeros((channels,))}
+        self.conv2 = {"weight": mx.zeros((channels, channels, 3, 3, 3)),
+                      "bias": mx.zeros((channels,))}
+
+    def __call__(self, x: mx.array) -> mx.array:
+        dtype = x.dtype
+        h = self.norm1(x).astype(dtype)
+        h = nn.silu(h)
+        h = conv3d(h, self.conv1["weight"], self.conv1["bias"])
+        h = self.norm2(h).astype(dtype)
+        h = nn.silu(h)
+        h = conv3d(h, self.conv2["weight"], self.conv2["bias"])
+        return h + x
+
+
+class PixelShuffleUpsample3d(nn.Module):
+    """Conv3d then 3D pixel shuffle (factor 2): channels/8 at 2x resolution."""
+
+    def __init__(self, in_channels: int, out_channels: int):
+        super().__init__()
+        # Output channels = out_channels * 8 (for 2^3 pixel shuffle)
+        self.out_channels = out_channels
+        self.conv = {"weight": mx.zeros((out_channels * 8, in_channels, 3, 3, 3)),
+                     "bias": mx.zeros((out_channels * 8,))}
+
+    def __call__(self, x: mx.array) -> mx.array:
+        h = conv3d(x, self.conv["weight"], self.conv["bias"])
+        # Pixel shuffle 3D: (B, D, H, W, C*8) → (B, D*2, H*2, W*2, C)
+        B, D, H, W, C8 = h.shape
+        C = self.out_channels
+        h = h.reshape(B, D, H, W, C, 2, 2, 2)
+        h = h.transpose(0, 1, 5, 2, 6, 3, 7, 4)  # (B, D, 2, H, 2, W, 2, C)
+        h = h.reshape(B, D * 2, H * 2, W * 2, C)
+        return h
+
+
+class MlxSparseStructureDecoder(nn.Module):
+    """
+    Dense 3D conv decoder for sparse structure prediction.
+    Input: (B, latent_channels, D, H, W) — typically (1, 8, 16, 16, 16)
+    Output: (B, 1, D', H', W') — binary occupancy grid
+    """
+
+    def __init__(self, out_channels: int = 1, latent_channels: int = 8,
+                 num_res_blocks: int = 2, num_res_blocks_middle: int = 2,
+                 channels: list = None, use_fp16: bool = True):
+        super().__init__()
+        if channels is None:
+            channels = [512, 128, 32]
+
+        self.use_fp16 = use_fp16
+
+        # Input layer
+        self.input_layer = {"weight": mx.zeros((channels[0], latent_channels, 3, 3, 3)),
+                            "bias": mx.zeros((channels[0],))}
+
+        # Middle blocks
+        self.middle_block = [ResBlock3d(channels[0]) for _ in range(num_res_blocks_middle)]
+
+        # Decoder blocks: [ResBlock × num_res_blocks, PixelShuffleUpsample] × stages
+        self.blocks = []
+        for i in range(len(channels)):
+            for _ in range(num_res_blocks):
+                self.blocks.append(ResBlock3d(channels[i]))
+            if i < len(channels) - 1:
+                self.blocks.append(PixelShuffleUpsample3d(channels[i], channels[i + 1]))
+
+        # Output layer: ChannelLayerNorm → SiLU → Conv3d
+        final_ch = channels[-1]
+        self.out_layer = [
+            ChannelLayerNorm3d(final_ch),
+            None,  # SiLU
+            {"weight": mx.zeros((out_channels, final_ch, 3, 3, 3)),
+             "bias": mx.zeros((out_channels,))},
+        ]
+
+        self._compiled_forward = None
+
+    def _run_blocks(self, h):
+        """Middle + decoder blocks — compilable dense graph."""
+        for block in self.middle_block:
+            h = block(h)
+        for block in self.blocks:
+            h = block(h)
+        return h
+
+    def __call__(self, x: mx.array) -> mx.array:
+        """x: (B, C, D, H, W) in PyTorch format → convert to (B, D, H, W, C) for MLX."""
+        # Convert from (B, C, D, H, W) to (B, D, H, W, C)
+        h = x.transpose(0, 2, 3, 4, 1)
+
+        if self.use_fp16:
+            h = h.astype(mx.float16)
+
+        # Input conv
+        h = conv3d(h, self.input_layer["weight"], self.input_layer["bias"])
+
+        # Middle + decoder blocks (compiled)
+        if self._compiled_forward is None:
+            try:
+                self._compiled_forward = mx.compile(self._run_blocks)
+            except Exception:
+                self._compiled_forward = False
+        if self._compiled_forward:
+            h = self._compiled_forward(h)
+        else:
+            h = self._run_blocks(h)
+
+        # Output
+        h = h.astype(mx.float32)
+        h = self.out_layer[0](h)  # GroupNorm
+        h = nn.silu(h)
+        h = conv3d(h, self.out_layer[2]["weight"], self.out_layer[2]["bias"])
+
+        # Convert back to (B, C, D, H, W)
+        h = h.transpose(0, 4, 1, 2, 3)
+        return h
+
+
+def load_structure_decoder(model_path: str) -> MlxSparseStructureDecoder:
+    """Load structure decoder from config + safetensors."""
+    import json
+
+    with open(f"{model_path}.json") as f:
+        config = json.load(f)
+
+    args = config['args']
+    model = MlxSparseStructureDecoder(
+        out_channels=args['out_channels'],
+        latent_channels=args['latent_channels'],
+        num_res_blocks=args['num_res_blocks'],
+        num_res_blocks_middle=args['num_res_blocks_middle'],
+        channels=args['channels'],
+        use_fp16=args.get('use_fp16', True),
+    )
+
+    weights = mx.load(f"{model_path}.safetensors")
+
+    # Remap weight keys for MLX module structure
+    remapped = _remap_structure_decoder_weights(weights)
+    model.load_weights(list(remapped.items()))
+    return model
+
+
+def _remap_structure_decoder_weights(weights: dict) -> dict:
+    """Remap PyTorch weight keys to MLX module paths."""
+    import re
+    remapped = {}
+    for k, v in weights.items():
+        new_k = k
+        # middle_block.N.xxx → middle_block.N.xxx (list indexing, no .layers.)
+        # blocks.N.xxx → blocks.N.xxx
+        # out_layer.0.xxx → out_layer.0.xxx (GroupNorm)
+        # out_layer.2.xxx → out_layer.2.xxx (Conv)
+
+        # For dict-based conv layers, map conv1.weight → conv1.weight etc.
+        # These are already in the right format for our dict-based storage
+        remapped[new_k] = v
+    return remapped

mlx_backend/transformer_block.py

@@ -0,0 +1,235 @@
+"""
+Modulated Sparse Transformer Cross Block for MLX.
+Implements AdaLN modulation → self-attn → cross-attn → FFN.
+"""
+import mlx.core as mx
+import mlx.nn as nn
+from .norm import LayerNorm32
+from .attention import MlxMultiHeadAttention
+
+
+class MlxSparseFeedForwardNet(nn.Module):
+    """FFN with GELU activation."""
+
+    def __init__(self, channels: int, mlp_ratio: float = 4.0):
+        super().__init__()
+        hidden = int(channels * mlp_ratio)
+        self.mlp = MlxSparseSequential(
+            nn.Linear(channels, hidden),
+            nn.GELU(approx="precise"),
+            nn.Linear(hidden, channels),
+        )
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.mlp(x)
+
+
+class MlxSparseSequential(nn.Module):
+    """Sequential container that passes feats through layers."""
+
+    def __init__(self, *layers):
+        super().__init__()
+        self.layers = list(layers)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        for layer in self.layers:
+            x = layer(x)
+        return x
+
+
+class MlxModulatedSparseTransformerCrossBlock(nn.Module):
+    """
+    Sparse transformer cross-attention block with AdaLN modulation.
+    Matches PyTorch ModulatedSparseTransformerCrossBlock.
+
+    With share_mod=True: uses per-block learnable `modulation` param + shared `mod` input.
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        ctx_channels: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        use_rope: bool = False,
+        rope_freq: tuple = (1.0, 10000.0),
+        share_mod: bool = False,
+        qk_rms_norm: bool = False,
+        qk_rms_norm_cross: bool = False,
+        qkv_bias: bool = True,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.share_mod = share_mod
+
+        self.norm1 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+        self.norm2 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
+        self.norm3 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+
+        self.self_attn = MlxMultiHeadAttention(
+            channels, num_heads,
+            type="self", qkv_bias=qkv_bias,
+            use_rope=use_rope, rope_freq=rope_freq,
+            qk_rms_norm=qk_rms_norm,
+        )
+        self.cross_attn = MlxMultiHeadAttention(
+            channels, num_heads,
+            ctx_channels=ctx_channels,
+            type="cross", qkv_bias=qkv_bias,
+            qk_rms_norm=qk_rms_norm_cross,
+        )
+        self.mlp = MlxSparseFeedForwardNet(channels, mlp_ratio=mlp_ratio)
+
+        if not share_mod:
+            self.adaLN_modulation = MlxSparseSequential(
+                nn.SiLU(),
+                nn.Linear(channels, 6 * channels, bias=True),
+            )
+        else:
+            self.modulation = mx.zeros((6 * channels,))
+
+    def __call__(
+        self,
+        x: mx.array,
+        mod: mx.array,
+        context: mx.array,
+        rope_cache: tuple = None,
+    ) -> mx.array:
+        """
+        Args:
+            x: (N, C) sparse features
+            mod: (1, 6*C) or (1, C) modulation signal from timestep
+            context: (M, ctx_C) conditioning features
+            rope_cache: (cos, sin) precomputed RoPE
+        """
+        if self.share_mod:
+            mods = (self.modulation + mod).astype(mod.dtype)
+        else:
+            mods = self.adaLN_modulation(mod)
+
+        # Split into 6 modulation signals: each (1, C)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = \
+            mx.split(mods, 6, axis=-1)
+
+        # Self-attention with AdaLN
+        h = self.norm1(x)
+        h = h * (1 + scale_msa) + shift_msa
+        h = self.self_attn(h, rope_cache=rope_cache)
+        h = h * gate_msa
+        x = x + h
+
+        # Cross-attention (no AdaLN, just norm2 with affine)
+        h = self.norm2(x)
+        h = self.cross_attn(h, context=context)
+        x = x + h
+
+        # FFN with AdaLN
+        h = self.norm3(x)
+        h = h * (1 + scale_mlp) + shift_mlp
+        h = self.mlp(h)
+        h = h * gate_mlp
+        x = x + h
+
+        return x
+
+
+class MlxModulatedTransformerCrossBlock(nn.Module):
+    """
+    Dense transformer cross-attention block with AdaLN modulation.
+    Used for SparseStructureFlowModel (dense 16^3 input).
+    Supports batched input (B, N, C) for batched CFG.
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        ctx_channels: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        use_rope: bool = False,
+        rope_freq: tuple = (1.0, 10000.0),
+        share_mod: bool = False,
+        qk_rms_norm: bool = False,
+        qk_rms_norm_cross: bool = False,
+        qkv_bias: bool = True,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.share_mod = share_mod
+
+        self.norm1 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+        self.norm2 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
+        self.norm3 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+
+        self.self_attn = MlxMultiHeadAttention(
+            channels, num_heads,
+            type="self", qkv_bias=qkv_bias,
+            use_rope=use_rope, rope_freq=rope_freq,
+            qk_rms_norm=qk_rms_norm,
+        )
+        self.cross_attn = MlxMultiHeadAttention(
+            channels, num_heads,
+            ctx_channels=ctx_channels,
+            type="cross", qkv_bias=qkv_bias,
+            qk_rms_norm=qk_rms_norm_cross,
+        )
+        self.mlp = MlxSparseFeedForwardNet(channels, mlp_ratio=mlp_ratio)
+
+        if not share_mod:
+            self.adaLN_modulation = MlxSparseSequential(
+                nn.SiLU(),
+                nn.Linear(channels, 6 * channels, bias=True),
+            )
+        else:
+            self.modulation = mx.zeros((6 * channels,))
+
+    def __call__(
+        self,
+        x: mx.array,
+        mod: mx.array,
+        context: mx.array,
+        rope_cache: tuple = None,
+    ) -> mx.array:
+        """
+        Args:
+            x: (N, C) or (B, N, C) dense features
+            mod: (B, 6*C) modulation signal
+            context: (M, ctx_C) or (B, M, ctx_C) conditioning features
+            rope_cache: (cos, sin) precomputed RoPE
+        """
+        if self.share_mod:
+            mods = (self.modulation + mod).astype(mod.dtype)
+        else:
+            mods = self.adaLN_modulation(mod)
+
+        if x.ndim == 3:
+            # Batched: mods is (B, 6*C), need (B, 1, C) for broadcasting
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = \
+                mx.split(mods, 6, axis=-1)
+            shift_msa = shift_msa[:, None, :]
+            scale_msa = scale_msa[:, None, :]
+            gate_msa = gate_msa[:, None, :]
+            shift_mlp = shift_mlp[:, None, :]
+            scale_mlp = scale_mlp[:, None, :]
+            gate_mlp = gate_mlp[:, None, :]
+        else:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = \
+                mx.split(mods, 6, axis=-1)
+
+        h = self.norm1(x)
+        h = h * (1 + scale_msa) + shift_msa
+        h = self.self_attn(h, rope_cache=rope_cache)
+        h = h * gate_msa
+        x = x + h
+
+        h = self.norm2(x)
+        h = self.cross_attn(h, context=context)
+        x = x + h
+
+        h = self.norm3(x)
+        h = h * (1 + scale_mlp) + shift_mlp
+        h = self.mlp(h)
+        h = h * gate_mlp
+        x = x + h
+
+        return x

mlx_backend/vae_blocks.py

@@ -0,0 +1,205 @@
+"""
+VAE decoder building blocks in MLX.
+Matches PyTorch SparseResBlock3d, SparseConvNeXtBlock3d, etc.
+"""
+import logging
+import time
+import mlx.core as mx
+import mlx.nn as nn
+from .norm import LayerNorm32
+from .sparse_tensor import MlxSparseTensor
+from .sparse_conv import MlxSparseConv3d
+from .sparse_ops import (
+    MlxSparseDownsample, MlxSparseUpsample,
+    MlxSparseSpatial2Channel, MlxSparseChannel2Spatial,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class MlxSparseLinear(nn.Module):
+    """Linear layer operating on sparse tensor features."""
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = True):
+        super().__init__()
+        self.linear = nn.Linear(in_features, out_features, bias=bias)
+
+    def __call__(self, x: MlxSparseTensor) -> MlxSparseTensor:
+        return x.replace(self.linear(x.feats))
+
+
+class MlxSparseConvNeXtBlock3d(nn.Module):
+    """ConvNeXt-style block: conv → norm → MLP residual."""
+
+    def __init__(self, channels: int, mlp_ratio: float = 4.0):
+        super().__init__()
+        self.channels = channels
+        self.norm = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
+        self.conv = MlxSparseConv3d(channels, channels, 3)
+        hidden = int(channels * mlp_ratio)
+        self.mlp = MlxConvNeXtMLP(channels, hidden)
+
+    def __call__(self, x: MlxSparseTensor) -> MlxSparseTensor:
+        t0 = time.time()
+        h = self.conv(x)
+        h = h.replace(self.norm(h.feats))
+        h = h.replace(self.mlp(h.feats))
+        result = MlxSparseTensor(
+            feats=h.feats + x.feats,
+            coords=x.coords,
+            shape=x.shape,
+            _scale=x._scale,
+            _spatial_cache=x._spatial_cache,
+        )
+        dt = time.time() - t0
+        if dt > 0.5:
+            logger.debug("[MLX] ConvNeXtBlock: N=%d, ch=%d, %.2fs",
+                         x.feats.shape[0], self.channels, dt)
+        return result
+
+
+class MlxConvNeXtMLP(nn.Module):
+    """MLP for ConvNeXt block: Linear → SiLU → Linear (zero-init)."""
+
+    def __init__(self, channels: int, hidden: int):
+        super().__init__()
+        # layers list used by weight remapping (mlp.layers.0/2)
+        self.layers = [
+            nn.Linear(channels, hidden),
+            None,  # SiLU (not a module)
+            nn.Linear(hidden, channels),
+        ]
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.layers[2](nn.silu(self.layers[0](x)))
+
+
+class MlxSparseResBlockC2S3d(nn.Module):
+    """
+    Residual block with Channel2Spatial upsampling.
+    Used in shape decoder to go from coarse to fine resolution.
+    """
+
+    def __init__(self, channels: int, out_channels: int = None,
+                 pred_subdiv: bool = True):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.pred_subdiv = pred_subdiv
+
+        self.norm1 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
+        self.norm2 = LayerNorm32(self.out_channels, elementwise_affine=False, eps=1e-6)
+        # conv1 outputs out_channels * 8 (spatial expansion)
+        self.conv1 = MlxSparseConv3d(channels, self.out_channels * 8, 3)
+        self.conv2 = MlxSparseConv3d(self.out_channels, self.out_channels, 3)
+        if pred_subdiv:
+            self.to_subdiv = MlxSparseLinear(channels, 8)
+        self.updown = MlxSparseChannel2Spatial(2)
+
+    def _skip_connection(self, x: MlxSparseTensor) -> MlxSparseTensor:
+        """Repeat features to match out_channels after C2S."""
+        # x after C2S has channels // 8 features
+        feats = x.feats
+        c_in = feats.shape[1]
+        repeats = self.out_channels // c_in
+        if repeats > 1:
+            feats = mx.repeat(feats, repeats, axis=1)
+        return x.replace(feats[:, :self.out_channels])
+
+    def __call__(self, x: MlxSparseTensor, subdiv: MlxSparseTensor = None):
+        t0 = time.time()
+        N_in = x.feats.shape[0]
+
+        if self.pred_subdiv:
+            subdiv = self.to_subdiv(x)
+
+        h = x.replace(self.norm1(x.feats))
+        h = h.replace(nn.silu(h.feats))
+        h = self.conv1(h)
+
+        subdiv_bin = subdiv.replace((subdiv.feats > 0).astype(subdiv.feats.dtype)) if subdiv is not None else None
+        h = self.updown(h, subdiv_bin)
+        x = self.updown(x, subdiv_bin)
+
+        h = h.replace(self.norm2(h.feats))
+        h = h.replace(nn.silu(h.feats))
+        h = self.conv2(h)
+
+        skip = self._skip_connection(x)
+        result = MlxSparseTensor(
+            feats=h.feats + skip.feats,
+            coords=h.coords,
+            shape=h.shape,
+            _scale=h._scale,
+            _spatial_cache=h._spatial_cache,
+        )
+
+        dt = time.time() - t0
+        logger.debug("[MLX] ResBlockC2S: N %d → %d, ch %d → %d, %.2fs",
+                     N_in, result.feats.shape[0], self.channels, self.out_channels, dt)
+
+        if self.pred_subdiv:
+            return result, subdiv
+        return result
+
+
+class MlxSparseResBlockUpsample3d(nn.Module):
+    """Residual block with nearest-neighbor upsampling."""
+
+    def __init__(self, channels: int, out_channels: int = None,
+                 pred_subdiv: bool = True):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.pred_subdiv = pred_subdiv
+
+        self.norm1 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
+        self.norm2 = LayerNorm32(self.out_channels, elementwise_affine=False, eps=1e-6)
+        self.conv1 = MlxSparseConv3d(channels, self.out_channels, 3)
+        self.conv2 = MlxSparseConv3d(self.out_channels, self.out_channels, 3)
+        if channels != self.out_channels:
+            self.skip_connection = MlxSparseLinear(channels, self.out_channels)
+        else:
+            self.skip_connection = None
+        if pred_subdiv:
+            self.to_subdiv = MlxSparseLinear(channels, 8)
+        self.updown = MlxSparseUpsample(2)
+
+    def __call__(self, x: MlxSparseTensor, subdiv: MlxSparseTensor = None):
+        t0 = time.time()
+        N_in = x.feats.shape[0]
+
+        if self.pred_subdiv:
+            subdiv = self.to_subdiv(x)
+
+        h = x.replace(self.norm1(x.feats))
+        h = h.replace(nn.silu(h.feats))
+        subdiv_bin = subdiv.replace((subdiv.feats > 0).astype(subdiv.feats.dtype)) if subdiv is not None else None
+        h = self.updown(h, subdiv_bin)
+        x_up = self.updown(x, subdiv_bin)
+
+        h = self.conv1(h)
+        h = h.replace(self.norm2(h.feats))
+        h = h.replace(nn.silu(h.feats))
+        h = self.conv2(h)
+
+        if self.skip_connection is not None:
+            skip = self.skip_connection(x_up)
+        else:
+            skip = x_up
+
+        result = MlxSparseTensor(
+            feats=h.feats + skip.feats,
+            coords=h.coords,
+            shape=h.shape,
+            _scale=h._scale,
+            _spatial_cache=h._spatial_cache,
+        )
+
+        dt = time.time() - t0
+        logger.debug("[MLX] ResBlockUpsample: N %d → %d, ch %d → %d, %.2fs",
+                     N_in, result.feats.shape[0], self.channels, self.out_channels, dt)
+
+        if self.pred_subdiv:
+            return result, subdiv
+        return result

mlx_backend/vae_decoders.py

@@ -0,0 +1,205 @@
+"""
+VAE decoders in MLX: SparseUnetVaeDecoder and FlexiDualGridVaeDecoder.
+"""
+import logging
+import time
+import mlx.core as mx
+import mlx.nn as nn
+from typing import List, Optional, Tuple
+from .norm import LayerNorm32
+from .sparse_tensor import MlxSparseTensor
+from .sparse_conv import MlxSparseConv3d
+from .vae_blocks import (
+    MlxSparseLinear,
+    MlxSparseConvNeXtBlock3d,
+    MlxSparseResBlockC2S3d,
+    MlxSparseResBlockUpsample3d,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class MlxSparseUnetVaeDecoder(nn.Module):
+    """
+    Sparse UNet VAE decoder.
+    Matches PyTorch SparseUnetVaeDecoder architecture.
+    """
+
+    def __init__(
+        self,
+        out_channels: int,
+        model_channels: list,
+        latent_channels: int,
+        num_blocks: list,
+        block_type: list,
+        up_block_type: list,
+        block_args: list = None,
+        use_fp16: bool = False,
+        pred_subdiv: bool = True,
+    ):
+        super().__init__()
+        self.out_channels = out_channels
+        self.model_channels = model_channels
+        self.pred_subdiv = pred_subdiv
+        self.use_fp16 = use_fp16
+
+        self.output_layer = MlxSparseLinear(model_channels[-1], out_channels)
+        self.from_latent = MlxSparseLinear(latent_channels, model_channels[0])
+
+        # Build blocks
+        BLOCK_TYPES = {
+            'SparseConvNeXtBlock3d': MlxSparseConvNeXtBlock3d,
+        }
+        UP_BLOCK_TYPES = {
+            'SparseResBlockC2S3d': MlxSparseResBlockC2S3d,
+            'SparseResBlockUpsample3d': MlxSparseResBlockUpsample3d,
+        }
+
+        self.blocks = []
+        for i in range(len(num_blocks)):
+            stage = []
+            BlockClass = BLOCK_TYPES[block_type[i]]
+            for j in range(num_blocks[i]):
+                stage.append(BlockClass(model_channels[i]))
+            if i < len(num_blocks) - 1:
+                UpBlockClass = UP_BLOCK_TYPES[up_block_type[i]]
+                stage.append(UpBlockClass(
+                    model_channels[i], model_channels[i + 1],
+                    pred_subdiv=pred_subdiv,
+                ))
+            self.blocks.append(stage)
+
+    def __call__(
+        self,
+        x: MlxSparseTensor,
+        guide_subs: list = None,
+        return_subs: bool = False,
+    ):
+        logger.info("[MLX] VAE decoder forward: N=%d, latent_ch=%d", x.feats.shape[0], x.feats.shape[1])
+        t_total = time.time()
+
+        h = self.from_latent(x)
+        if self.use_fp16:
+            h = h.replace(h.feats.astype(mx.float16))
+
+        subs = []
+        for i, stage in enumerate(self.blocks):
+            t_stage = time.time()
+            for j, block in enumerate(stage):
+                if i < len(self.blocks) - 1 and j == len(stage) - 1:
+                    # Up block
+                    if self.pred_subdiv:
+                        h, sub = block(h)
+                        subs.append(sub)
+                    else:
+                        h = block(h, subdiv=guide_subs[i] if guide_subs is not None else None)
+                else:
+                    h = block(h)
+
+            # Eval once per stage (not per block) to bound memory
+            mx.eval(h.feats)
+            dt_stage = time.time() - t_stage
+            logger.info("[MLX] VAE stage %d/%d: N=%d, channels=%d, %.2fs",
+                        i + 1, len(self.blocks), h.feats.shape[0], h.feats.shape[1], dt_stage)
+
+        h = h.replace(h.feats.astype(x.feats.dtype))
+
+        # Final layer norm (two-pass for parity)
+        feats = h.feats.astype(mx.float32)
+        mean = mx.mean(feats, axis=-1, keepdims=True)
+        centered = feats - mean
+        var = mx.mean(centered * centered, axis=-1, keepdims=True)
+        feats = centered * mx.rsqrt(var + 1e-5)
+        h = h.replace(feats.astype(x.feats.dtype))
+
+        h = self.output_layer(h)
+
+        dt_total = time.time() - t_total
+        logger.info("[MLX] VAE decoder done: N=%d, out_ch=%d, total %.2fs",
+                    h.feats.shape[0], h.feats.shape[1], dt_total)
+
+        if return_subs:
+            return h, subs
+        return h
+
+    def upsample(self, x: MlxSparseTensor, upsample_times: int) -> mx.array:
+        """Run decoder up to upsample_times stages, return upsampled coords."""
+        logger.info("[MLX] VAE upsample: N=%d, upsample_times=%d", x.feats.shape[0], upsample_times)
+        h = self.from_latent(x)
+        if self.use_fp16:
+            h = h.replace(h.feats.astype(mx.float16))
+
+        for i, stage in enumerate(self.blocks):
+            if i == upsample_times:
+                return h.coords
+            for j, block in enumerate(stage):
+                if i < len(self.blocks) - 1 and j == len(stage) - 1:
+                    h, sub = block(h)
+                else:
+                    h = block(h)
+                mx.eval(h.feats)
+
+        return h.coords
+
+
+class MlxFlexiDualGridVaeDecoder(nn.Module):
+    """
+    FlexiDualGrid VAE decoder — wraps SparseUnetVaeDecoder.
+    Outputs mesh vertices + intersection logits + quad_lerp.
+    """
+
+    def __init__(
+        self,
+        resolution: int,
+        model_channels: list,
+        latent_channels: int,
+        num_blocks: list,
+        block_type: list,
+        up_block_type: list,
+        block_args: list = None,
+        use_fp16: bool = False,
+        voxel_margin: float = 0.5,
+    ):
+        super().__init__()
+        self.resolution = resolution
+        self.voxel_margin = voxel_margin
+
+        # out_channels = 7 (3 vertex + 3 intersection + 1 quad_lerp)
+        self.decoder = MlxSparseUnetVaeDecoder(
+            out_channels=7,
+            model_channels=model_channels,
+            latent_channels=latent_channels,
+            num_blocks=num_blocks,
+            block_type=block_type,
+            up_block_type=up_block_type,
+            block_args=block_args,
+            use_fp16=use_fp16,
+            pred_subdiv=True,
+        )
+
+    def set_resolution(self, resolution: int):
+        self.resolution = resolution
+
+    def __call__(self, x: MlxSparseTensor, return_subs: bool = False, **kwargs):
+        decoded = self.decoder(x, return_subs=return_subs, **kwargs)
+
+        if return_subs:
+            h, subs = decoded
+        else:
+            h = decoded
+            subs = None
+
+        # Post-process: vertices = sigmoid(h[:, :3]), intersected = h[:, 3:6] > 0
+        feats = h.feats
+        vertices_feats = (1 + 2 * self.voxel_margin) * mx.sigmoid(feats[:, :3]) - self.voxel_margin
+        intersected_feats = (feats[:, 3:6] > 0).astype(feats.dtype)
+        quad_lerp_feats = mx.where(feats[:, 6:7] > 0, feats[:, 6:7], mx.zeros_like(feats[:, 6:7])) + mx.log1p(mx.exp(-mx.abs(feats[:, 6:7])))  # softplus
+
+        result = (h, vertices_feats, intersected_feats, quad_lerp_feats)
+
+        if return_subs:
+            return result, subs
+        return result
+
+    def upsample(self, x: MlxSparseTensor, upsample_times: int) -> mx.array:
+        return self.decoder.upsample(x, upsample_times)

o-voxel/o_voxel/__init__.py

@@ -1,7 +1,17 @@
 from . import (
     convert,
-    io,
     postprocess,
-    rasterize,
-    serialize
+    postprocess_cpu,
 )
+
+try:
+    from . import io
+except ImportError:
+    # io.vxz requires serialize which requires _C extension
+    pass
+
+try:
+    from . import rasterize, serialize
+except ImportError:
+    # rasterize and serialize require CUDA _C extension
+    pass

o-voxel/o_voxel/convert/__init__.py

@@ -1,2 +1,6 @@
 from .flexible_dual_grid import *
-from .volumetic_attr import *
+try:
+    from .volumetic_attr import *
+except ImportError:
+    # volumetic_attr requires _C extension
+    pass

o-voxel/o_voxel/convert/flexible_dual_grid.py

@@ -1,7 +1,14 @@
 from typing import *
 import numpy as np
 import torch
-from .. import _C
+import platform
+
+_HAS_C = False
+try:
+    from .. import _C
+    _HAS_C = True
+except ImportError:
+    pass
 
 __all__ = [
     "mesh_to_flexible_dual_grid",
@@ -25,6 +32,44 @@ def _init_hashmap(grid_size, capacity, device):
     return hashmap_keys, hashmap_vals
 
 
+class _CPUHashMap:
+    """Pure PyTorch hashmap replacement for CUDA _C.hashmap_* functions."""
+
+    def __init__(self, grid_size, device='cpu'):
+        D, H, W = int(grid_size[0]), int(grid_size[1]), int(grid_size[2])
+        self.D, self.H, self.W = D, H, W
+        self.table_size = D * H * W
+        self.device = device
+        # Use int64 flat lookup table
+        self.lookup = torch.full((self.table_size,), -1, dtype=torch.long, device=device)
+
+    def _flat_key(self, coords_3d):
+        """coords_3d: (..., 3) int tensor of (x, y, z)"""
+        return (coords_3d[..., 0].long() * self.H * self.W +
+                coords_3d[..., 1].long() * self.W +
+                coords_3d[..., 2].long())
+
+    def insert(self, coords_4d):
+        """coords_4d: (N, 4) with [batch, x, y, z]. batch is ignored (assumed 0), value = row index."""
+        flat = self._flat_key(coords_4d[:, 1:4])
+        self.lookup[flat] = torch.arange(coords_4d.shape[0], dtype=torch.long, device=self.device)
+
+    def lookup_3d(self, coords_4d):
+        """coords_4d: (M, 4) with [batch, x, y, z]. Returns (M,) indices, 0xffffffff for missing."""
+        coords_3d = coords_4d[:, 1:4]
+        flat = self._flat_key(coords_3d)
+        # Bounds check
+        valid = ((coords_3d[..., 0] >= 0) & (coords_3d[..., 0] < self.D) &
+                 (coords_3d[..., 1] >= 0) & (coords_3d[..., 1] < self.H) &
+                 (coords_3d[..., 2] >= 0) & (coords_3d[..., 2] < self.W))
+        flat = flat.clamp(0, self.table_size - 1)
+        result = self.lookup[flat]
+        result[~valid] = -1
+        # Convert -1 to 0xffffffff for compatibility
+        result[result < 0] = 0xffffffff
+        return result
+
+
 @torch.no_grad()
 def mesh_to_flexible_dual_grid(
     vertices: torch.Tensor,
@@ -113,7 +158,7 @@ def mesh_to_flexible_dual_grid(
             min_xyz -= padding * 0.5
             max_xyz += padding * 0.5
 
-        aabb = torch.stack([min_xyz, max_xyz], dim=0).float().cuda()
+        aabb = torch.stack([min_xyz, max_xyz], dim=0).float().to(vertices.device)
 
     # Fill voxel size or grid size
     if voxel_size is None:
@@ -222,8 +267,14 @@ def flexible_dual_grid_to_mesh(
     mesh_vertices = (coords.float() + dual_vertices) / (2 * N) - 0.5
 
     # Store active voxels into hashmap
+    if _HAS_C and coords.is_cuda:
         hashmap = _init_hashmap(grid_size, 2 * N, device=coords.device)
         _C.hashmap_insert_3d_idx_as_val_cuda(*hashmap, torch.cat([torch.zeros_like(coords[:, :1]), coords], dim=-1), *grid_size.tolist())
+        _use_cpu_hashmap = False
+    else:
+        cpu_hashmap = _CPUHashMap(grid_size, device=coords.device)
+        cpu_hashmap.insert(torch.cat([torch.zeros_like(coords[:, :1]), coords], dim=-1))
+        _use_cpu_hashmap = True
 
     # Find connected voxels
     edge_neighbor_voxel = coords.reshape(N, 1, 1, 3) + flexible_dual_grid_to_mesh.edge_neighbor_voxel_offset      # (N, 3, 4, 3)
@@ -233,6 +284,9 @@ def flexible_dual_grid_to_mesh(
         torch.zeros((M * 4, 1), dtype=torch.int, device=coords.device),
         connected_voxel.reshape(-1, 3)
     ], dim=1)
+    if _use_cpu_hashmap:
+        connected_voxel_indices = cpu_hashmap.lookup_3d(connected_voxel_hash_key).reshape(M, 4).int()
+    else:
         connected_voxel_indices = _C.hashmap_lookup_3d_cuda(*hashmap, connected_voxel_hash_key, *grid_size.tolist()).reshape(M, 4).int()
     connected_voxel_valid = (connected_voxel_indices != 0xffffffff).all(dim=1)
     quad_indices = connected_voxel_indices[connected_voxel_valid].int()                             # (L, 4)
@@ -277,7 +331,7 @@ def flexible_dual_grid_to_mesh(
             split_weight_ws_13 * mean_v13
         ) / (split_weight_ws_02 + split_weight_ws_13)
         mesh_vertices = torch.cat([mesh_vertices, mid_vertices], dim=0)
-        quad_indices = torch.cat([quad_indices, torch.arange(N, N + L, device='cuda').unsqueeze(1)], dim=1)
+        quad_indices = torch.cat([quad_indices, torch.arange(N, N + L, device=coords.device).unsqueeze(1)], dim=1)
         mesh_triangles = quad_indices[:, flexible_dual_grid_to_mesh.quad_split_train].reshape(-1, 3)
     
     return mesh_vertices, mesh_triangles

o-voxel/o_voxel/postprocess.py

@@ -6,9 +6,69 @@ import cv2
 from PIL import Image
 import trimesh
 import trimesh.visual
-from flex_gemm.ops.grid_sample import grid_sample_3d
-import nvdiffrast.torch as dr
-import cumesh
+
+import platform
+
+_HAS_DR = False
+_HAS_MESH = False
+_BACKEND = None
+dr = None
+
+# Differentiable rasterizer — mtldiffrast (Metal) or nvdiffrast (CUDA)
+try:
+    import mtldiffrast.torch as dr
+    _HAS_DR = True
+    _BACKEND = 'metal'
+except ImportError:
+    try:
+        import nvdiffrast.torch as dr
+        _HAS_DR = True
+        _BACKEND = 'cuda'
+    except ImportError:
+        pass
+
+# Mesh processing — cumesh auto-selects Metal/CUDA
+try:
+    import cumesh
+    _MeshBackend = cumesh.CuMesh
+    _BVH = cumesh.cuBVH
+    _remesh_narrow_band_dc = cumesh.remeshing.remesh_narrow_band_dc
+    _HAS_MESH = True
+    if _BACKEND is None:
+        _BACKEND = 'metal' if platform.system() == 'Darwin' else 'cuda'
+except ImportError:
+    pass
+
+_HAS_GPU_DEPS = _HAS_DR and _HAS_MESH
+
+try:
+    from flex_gemm.ops.grid_sample import grid_sample_3d as _flex_grid_sample_3d
+    _HAS_FLEX_GEMM = True
+except ImportError:
+    _HAS_FLEX_GEMM = False
+
+
+def _grid_sample_3d(feats, coords, shape, grid, mode='trilinear'):
+    """Grid sampling with flex_gemm on CUDA, F.grid_sample fallback otherwise."""
+    if _HAS_FLEX_GEMM:
+        return _flex_grid_sample_3d(feats, coords, shape, grid, mode=mode)
+    import torch.nn.functional as F_gs
+    B, C = shape[0], shape[1]
+    D, H, W = shape[2], shape[3], shape[4]
+    device = feats.device
+    dense_vol = torch.zeros(B, C, D, H, W, dtype=feats.dtype, device=device)
+    batch_idx = coords[:, 0].long()
+    cx, cy, cz = coords[:, 1].long(), coords[:, 2].long(), coords[:, 3].long()
+    dense_vol[batch_idx, :, cx, cy, cz] = feats
+    grid_norm = torch.stack([
+        grid[..., 2] / (W - 1) * 2 - 1,
+        grid[..., 1] / (H - 1) * 2 - 1,
+        grid[..., 0] / (D - 1) * 2 - 1,
+    ], dim=-1).reshape(B, 1, 1, -1, 3)
+    sampled = F_gs.grid_sample(dense_vol, grid_norm, mode='bilinear',
+                               align_corners=True, padding_mode='border')
+    M = grid.shape[1]
+    return sampled.reshape(B * C, M)
 
 
 def to_glb(
@@ -57,6 +117,27 @@ def to_glb(
         verbose: whether to print verbose messages
         use_tqdm: whether to use tqdm to display progress bar
     """
+    # Auto-fallback to CPU pipeline when no GPU deps available
+    if not _HAS_GPU_DEPS:
+        from .postprocess_cpu import to_glb as to_glb_cpu
+        return to_glb_cpu(
+            vertices=vertices, faces=faces, attr_volume=attr_volume,
+            coords=coords, attr_layout=attr_layout, aabb=aabb,
+            voxel_size=voxel_size, grid_size=grid_size,
+            decimation_target=decimation_target, texture_size=texture_size,
+            remesh=remesh, remesh_band=remesh_band, remesh_project=remesh_project,
+            verbose=verbose, use_tqdm=use_tqdm,
+        )
+
+    # Select device based on backend
+    # Metal path: all GPU compute goes through Metal kernels directly (mtldiffrast,
+    # mtlbvh, cumesh, flex_gemm). CPU tensors on Apple Silicon unified memory are
+    # directly GPU-accessible — no MPS overhead needed.
+    if _BACKEND == 'metal':
+        device = torch.device('cpu')
+    else:
+        device = torch.device('cuda')
+
     # --- Input Normalization (AABB, Voxel Size, Grid Size) ---
     if isinstance(aabb, (list, tuple)):
         aabb = np.array(aabb)
@@ -98,11 +179,11 @@ def to_glb(
         print(f"Original mesh: {vertices.shape[0]} vertices, {faces.shape[0]} faces")
 
     # Move data to GPU
-    vertices = vertices.cuda()
-    faces = faces.cuda()
+    vertices = vertices.to(device)
+    faces = faces.to(device)
 
-    # Initialize CUDA mesh handler
-    mesh = cumesh.CuMesh()
+    # Initialize mesh handler
+    mesh = _MeshBackend()
     mesh.init(vertices, faces)
     
     # --- Initial Mesh Cleaning ---
@@ -119,7 +200,7 @@ def to_glb(
         pbar.set_description("Building BVH")
     if verbose:
         print(f"Building BVH for current mesh...", end='', flush=True)
-    bvh = cumesh.cuBVH(vertices, faces)
+    bvh = _BVH(vertices, faces)
     if use_tqdm:
         pbar.update(1)
     if verbose:
@@ -168,7 +249,7 @@ def to_glb(
         resolution = grid_size.max().item()
         
         # Perform Dual Contouring remeshing (rebuilds topology)
-        mesh.init(*cumesh.remeshing.remesh_narrow_band_dc(
+        mesh.init(*_remesh_narrow_band_dc(
             vertices, faces,
             center = center,
             scale = (resolution + 3 * remesh_band) / resolution * scale,
@@ -181,7 +262,15 @@ def to_glb(
         if verbose:
             print(f"After remeshing: {mesh.num_vertices} vertices, {mesh.num_faces} faces")
 
-        # Simplify and clean the remeshed result (similar logic to above)
+        # Clean up topology before simplification
+        mesh.remove_duplicate_faces()
+        mesh.repair_non_manifold_edges()
+        mesh.remove_small_connected_components(1e-5)
+        mesh.fill_holes(max_hole_perimeter=3e-2)
+        if verbose:
+            print(f"After cleanup: {mesh.num_vertices} vertices, {mesh.num_faces} faces")
+
+        # Simplify and clean the remeshed result
         mesh.simplify(decimation_target, verbose=verbose)
         if verbose:
             print(f"After simplifying: {mesh.num_vertices} vertices, {mesh.num_faces} faces")
@@ -208,10 +297,10 @@ def to_glb(
         return_vmaps=True,
         verbose=verbose,
     )
-    out_vertices = out_vertices.cuda()
-    out_faces = out_faces.cuda()
-    out_uvs = out_uvs.cuda()
-    out_vmaps = out_vmaps.cuda()
+    out_vertices = out_vertices.to(device)
+    out_faces = out_faces.to(device)
+    out_uvs = out_uvs.to(device)
+    out_vmaps = out_vmaps.to(device)
     mesh.compute_vertex_normals()
     out_normals = mesh.read_vertex_normals()[out_vmaps]
     
@@ -227,10 +316,10 @@ def to_glb(
         print("Sampling attributes...", end='', flush=True)
         
     # Setup differentiable rasterizer context
-    ctx = dr.RasterizeCudaContext()
+    ctx = dr.MtlRasterizeContext() if _BACKEND == 'metal' else dr.RasterizeCudaContext()
     # Prepare UV coordinates for rasterization (rendering in UV space)
     uvs_rast = torch.cat([out_uvs * 2 - 1, torch.zeros_like(out_uvs[:, :1]), torch.ones_like(out_uvs[:, :1])], dim=-1).unsqueeze(0)
-    rast = torch.zeros((1, texture_size, texture_size, 4), device='cuda', dtype=torch.float32)
+    rast = torch.zeros((1, texture_size, texture_size, 4), device=device, dtype=torch.float32)
     
     # Rasterize in chunks to save memory
     for i in range(0, out_faces.shape[0], 100000):
@@ -256,8 +345,8 @@ def to_glb(
     valid_pos = (orig_tri_verts * uvw.unsqueeze(-1)).sum(dim=1)
     
     # Trilinear sampling from the attribute volume (Color, Material props)
-    attrs = torch.zeros(texture_size, texture_size, attr_volume.shape[1], device='cuda')
-    attrs[mask] = grid_sample_3d(
+    attrs = torch.zeros(texture_size, texture_size, attr_volume.shape[1], device=device)
+    attrs[mask] = _grid_sample_3d(
         attr_volume,
         torch.cat([torch.zeros_like(coords[:, :1]), coords], dim=-1),
         shape=torch.Size([1, attr_volume.shape[1], *grid_size.tolist()]),
@@ -282,6 +371,13 @@ def to_glb(
     metallic = np.clip(attrs[..., attr_layout['metallic']].cpu().numpy() * 255, 0, 255).astype(np.uint8)
     roughness = np.clip(attrs[..., attr_layout['roughness']].cpu().numpy() * 255, 0, 255).astype(np.uint8)
     alpha = np.clip(attrs[..., attr_layout['alpha']].cpu().numpy() * 255, 0, 255).astype(np.uint8)
+    # Auto-detect transparency from baked alpha values
+    alpha_valid = alpha[mask]
+    if alpha_valid.size > 0 and alpha_valid.min() < 250:
+        alpha_mode = 'BLEND'
+        if verbose:
+            print(f"Detected transparency (alpha min={alpha_valid.min()}), using BLEND mode")
+    else:
         alpha_mode = 'OPAQUE'
     
     # Inpainting: fill gaps (dilation) to prevent black seams at UV boundaries
@@ -309,10 +405,10 @@ def to_glb(
     uvs_np = out_uvs.cpu().numpy()
     normals_np = out_normals.cpu().numpy()
     
-    # Swap Y and Z axes, invert Y (common conversion for GLB compatibility)
+    # Y-up to Z-up for GLB
     vertices_np[:, 1], vertices_np[:, 2] = vertices_np[:, 2], -vertices_np[:, 1]
     normals_np[:, 1], normals_np[:, 2] = normals_np[:, 2], -normals_np[:, 1]
-    uvs_np[:, 1] = 1 - uvs_np[:, 1] # Flip UV V-coordinate
+    uvs_np[:, 1] = 1 - uvs_np[:, 1]
     
     textured_mesh = trimesh.Trimesh(
         vertices=vertices_np,

o-voxel/o_voxel/postprocess_cpu.py

@@ -0,0 +1,419 @@
+"""
+macOS GLB export pipeline.
+
+Replaces cumesh + nvdiffrast + flex_gemm with:
+- fast_simplification / trimesh for mesh decimation
+- xatlas for UV unwrapping
+- PyTorch MPS-accelerated UV rasterization + texture baking
+- OpenCV cv2.inpaint for texture inpainting
+"""
+from typing import *
+from tqdm import tqdm
+import numpy as np
+import torch
+import torch.nn.functional as F
+import cv2
+from PIL import Image
+import trimesh
+import trimesh.visual
+import xatlas
+
+try:
+    import fast_simplification
+    HAS_FAST_SIMPLIFICATION = True
+except ImportError:
+    HAS_FAST_SIMPLIFICATION = False
+
+
+def _get_device():
+    """Get best available device: MPS > CPU."""
+    if torch.backends.mps.is_available():
+        return torch.device('mps')
+    return torch.device('cpu')
+
+
+def _rasterize_uv_gpu(vertices, faces, uvs, texture_size, device=None):
+    """
+    GPU-accelerated UV-space rasterization via vectorized PyTorch.
+
+    For each face, computes barycentric coords for all texels in its bounding box
+    in parallel on MPS/CPU. Replaces the slow Python double-loop.
+
+    Args:
+        vertices: (V, 3) vertex positions
+        faces: (F, 3) face indices (int64)
+        uvs: (V, 2) per-vertex UV coords in [0, 1]
+        texture_size: int
+        device: torch device (defaults to MPS if available)
+
+    Returns:
+        pos: (H, W, 3) interpolated 3D positions
+        mask: (H, W) bool - valid texels
+    """
+    if device is None:
+        device = _get_device()
+
+    H = W = texture_size
+
+    # Move to device
+    verts = vertices.float().to(device)
+    faces_t = faces.long().to(device)
+    uv = uvs.float().to(device)
+
+    # Gather per-face data: (F, 3, 2) UVs and (F, 3, 3) positions
+    face_uvs = uv[faces_t]      # (F, 3, 2)
+    face_verts = verts[faces_t]  # (F, 3, 3)
+
+    # Scale UVs to pixel coords
+    face_uvs_px = face_uvs.clone()
+    face_uvs_px[..., 0] *= (W - 1)
+    face_uvs_px[..., 1] *= (H - 1)
+
+    # Output buffers
+    pos_buf = torch.zeros(H, W, 3, device=device)
+    mask_buf = torch.zeros(H, W, dtype=torch.bool, device=device)
+    # Use a depth buffer (face index) to handle overlaps — last writer wins like nvdiffrast
+
+    # Process in chunks to avoid OOM
+    num_faces = faces_t.shape[0]
+    chunk_size = 50000
+
+    for start in range(0, num_faces, chunk_size):
+        end = min(start + chunk_size, num_faces)
+        chunk_uvs = face_uvs_px[start:end]   # (C, 3, 2)
+        chunk_verts = face_verts[start:end]   # (C, 3, 3)
+        C = chunk_uvs.shape[0]
+
+        # Bounding boxes per face: (C,)
+        bb_min_x = chunk_uvs[..., 0].min(dim=1).values.floor().clamp(min=0).int()
+        bb_max_x = chunk_uvs[..., 0].max(dim=1).values.ceil().clamp(max=W - 1).int()
+        bb_min_y = chunk_uvs[..., 1].min(dim=1).values.floor().clamp(min=0).int()
+        bb_max_y = chunk_uvs[..., 1].max(dim=1).values.ceil().clamp(max=H - 1).int()
+
+        # Compute max bbox size in this chunk to allocate grid
+        max_w = (bb_max_x - bb_min_x + 1).max().item()
+        max_h = (bb_max_y - bb_min_y + 1).max().item()
+
+        if max_w <= 0 or max_h <= 0:
+            continue
+
+        # Create local pixel grids for each face: (C, max_h, max_w, 2)
+        # Use offsets from bb_min
+        local_y = torch.arange(max_h, device=device).float()
+        local_x = torch.arange(max_w, device=device).float()
+        gy, gx = torch.meshgrid(local_y, local_x, indexing='ij')  # (max_h, max_w)
+        grid = torch.stack([gx, gy], dim=-1)  # (max_h, max_w, 2)
+        grid = grid.unsqueeze(0).expand(C, -1, -1, -1).clone()  # (C, max_h, max_w, 2)
+
+        # Offset to absolute pixel coords
+        grid[..., 0] += bb_min_x.float().view(C, 1, 1)
+        grid[..., 1] += bb_min_y.float().view(C, 1, 1)
+
+        # Validity mask: within bounding box
+        valid = (
+            (grid[..., 0] <= bb_max_x.float().view(C, 1, 1)) &
+            (grid[..., 1] <= bb_max_y.float().view(C, 1, 1)) &
+            (grid[..., 0] >= 0) & (grid[..., 0] < W) &
+            (grid[..., 1] >= 0) & (grid[..., 1] < H)
+        )  # (C, max_h, max_w)
+
+        # Compute barycentric coordinates for each texel relative to each face
+        # v0 = uv1 - uv0, v1 = uv2 - uv0, v2 = p - uv0
+        uv0 = chunk_uvs[:, 0, :]  # (C, 2)
+        v0 = chunk_uvs[:, 1, :] - uv0  # (C, 2)
+        v1 = chunk_uvs[:, 2, :] - uv0  # (C, 2)
+
+        # (C, max_h, max_w, 2)
+        v2 = grid - uv0.view(C, 1, 1, 2)
+
+        # Barycentric via dot products — all vectorized
+        dot00 = (v0 * v0).sum(dim=1)  # (C,)
+        dot01 = (v0 * v1).sum(dim=1)  # (C,)
+        dot11 = (v1 * v1).sum(dim=1)  # (C,)
+
+        dot02 = (v2 * v0.view(C, 1, 1, 2)).sum(dim=-1)  # (C, max_h, max_w)
+        dot12 = (v2 * v1.view(C, 1, 1, 2)).sum(dim=-1)  # (C, max_h, max_w)
+
+        inv_denom = dot00 * dot11 - dot01 * dot01  # (C,)
+        # Skip degenerate triangles
+        non_degen = inv_denom.abs() > 1e-10
+        inv_denom = torch.where(non_degen, 1.0 / inv_denom.clamp(min=1e-10), torch.zeros_like(inv_denom))
+
+        u = (dot11.view(C, 1, 1) * dot02 - dot01.view(C, 1, 1) * dot12) * inv_denom.view(C, 1, 1)
+        v = (dot00.view(C, 1, 1) * dot12 - dot01.view(C, 1, 1) * dot02) * inv_denom.view(C, 1, 1)
+        w = 1.0 - u - v
+
+        # Inside triangle test (with small epsilon for edge cases)
+        eps = -1e-4
+        inside = (u >= eps) & (v >= eps) & ((u + v) <= 1.0 - eps) & valid & non_degen.view(C, 1, 1)
+
+        # Interpolate 3D positions: w*v0 + u*v1 + v*v2
+        # chunk_verts: (C, 3, 3) — positions of 3 verts per face
+        interp_pos = (
+            w.unsqueeze(-1) * chunk_verts[:, 0:1, :].unsqueeze(2) +  # (C, 1, 1, 3)
+            u.unsqueeze(-1) * chunk_verts[:, 1:2, :].unsqueeze(2) +
+            v.unsqueeze(-1) * chunk_verts[:, 2:3, :].unsqueeze(2)
+        )  # (C, max_h, max_w, 3)
+
+        # Write to output buffer
+        # Get absolute pixel coords for valid texels
+        abs_x = grid[..., 0].long()
+        abs_y = grid[..., 1].long()
+
+        # Flatten and filter
+        inside_flat = inside.reshape(-1)
+        abs_x_flat = abs_x.reshape(-1)[inside_flat]
+        abs_y_flat = abs_y.reshape(-1)[inside_flat]
+        pos_flat = interp_pos.reshape(-1, 3)[inside_flat]
+
+        if abs_x_flat.numel() > 0:
+            pos_buf[abs_y_flat, abs_x_flat] = pos_flat
+            mask_buf[abs_y_flat, abs_x_flat] = True
+
+    return pos_buf.cpu().numpy(), mask_buf.cpu().numpy()
+
+
+def _grid_sample_3d_gpu(attr_volume, coords, grid_size, query_pts, voxel_size, aabb, device=None):
+    """
+    GPU-accelerated trilinear sampling from sparse attribute volume.
+    Uses F.grid_sample on MPS.
+    """
+    if device is None:
+        device = _get_device()
+
+    C = attr_volume.shape[1]
+    D, H, W = int(grid_size[0]), int(grid_size[1]), int(grid_size[2])
+
+    # Build dense volume on device
+    attr_vol = attr_volume.float().to(device)
+    coords_d = coords.long().to(device)
+
+    dense_vol = torch.zeros(1, C, D, H, W, dtype=torch.float32, device=device)
+    cx, cy, cz = coords_d[:, 0], coords_d[:, 1], coords_d[:, 2]
+    dense_vol[0, :, cx, cy, cz] = attr_vol.T
+
+    # Normalize query pts to [-1, 1] for grid_sample
+    query = query_pts.float().to(device)
+    aabb_d = aabb.float().to(device)
+    vs_d = voxel_size.float().to(device)
+
+    grid_pts = (query - aabb_d[0]) / vs_d
+    grid_pts_norm = torch.stack([
+        grid_pts[:, 2] / (W - 1) * 2 - 1,
+        grid_pts[:, 1] / (H - 1) * 2 - 1,
+        grid_pts[:, 0] / (D - 1) * 2 - 1,
+    ], dim=-1)
+    grid_pts_norm = grid_pts_norm.reshape(1, 1, 1, -1, 3)
+
+    sampled = F.grid_sample(dense_vol, grid_pts_norm, mode='bilinear', align_corners=True, padding_mode='border')
+    return sampled.reshape(C, -1).T.cpu()
+
+
+def to_glb(
+    vertices: torch.Tensor,
+    faces: torch.Tensor,
+    attr_volume: torch.Tensor,
+    coords: torch.Tensor,
+    attr_layout: Dict[str, slice],
+    aabb: Union[list, tuple, np.ndarray, torch.Tensor],
+    voxel_size: Union[float, list, tuple, np.ndarray, torch.Tensor] = None,
+    grid_size: Union[int, list, tuple, np.ndarray, torch.Tensor] = None,
+    decimation_target: int = 1000000,
+    texture_size: int = 2048,
+    remesh: bool = False,
+    remesh_band: float = 1,
+    remesh_project: float = 0.9,
+    verbose: bool = False,
+    use_tqdm: bool = False,
+    **kwargs,
+):
+    """
+    macOS GLB export. Replaces the CUDA pipeline (cumesh + nvdiffrast + flex_gemm)
+    with fast_simplification + xatlas + PyTorch MPS rasterization.
+    """
+    device = _get_device()
+    if verbose:
+        print(f"Using device: {device}")
+
+    # --- Input Normalization ---
+    if isinstance(aabb, (list, tuple)):
+        aabb = np.array(aabb)
+    if isinstance(aabb, np.ndarray):
+        aabb = torch.tensor(aabb, dtype=torch.float32)
+    aabb = aabb.cpu()
+
+    if voxel_size is not None:
+        if isinstance(voxel_size, (int, float)):
+            voxel_size = [voxel_size] * 3
+        if isinstance(voxel_size, (list, tuple)):
+            voxel_size = np.array(voxel_size)
+        if isinstance(voxel_size, np.ndarray):
+            voxel_size = torch.tensor(voxel_size, dtype=torch.float32)
+        grid_size = ((aabb[1] - aabb[0]) / voxel_size).round().int()
+    else:
+        assert grid_size is not None
+        if isinstance(grid_size, int):
+            grid_size = [grid_size] * 3
+        if isinstance(grid_size, (list, tuple)):
+            grid_size = np.array(grid_size)
+        if isinstance(grid_size, np.ndarray):
+            grid_size = torch.tensor(grid_size, dtype=torch.int32)
+        voxel_size = (aabb[1] - aabb[0]) / grid_size.float()
+
+    if use_tqdm:
+        pbar = tqdm(total=5, desc="Extracting GLB")
+    if verbose:
+        print(f"Original mesh: {vertices.shape[0]} vertices, {faces.shape[0]} faces")
+
+    vertices = vertices.cpu()
+    faces = faces.cpu()
+
+    # --- Step 1: Mesh Cleaning ---
+    if use_tqdm:
+        pbar.set_description("Cleaning mesh")
+    tm = trimesh.Trimesh(
+        vertices=vertices.numpy(),
+        faces=faces.numpy(),
+        process=False,
+    )
+    trimesh.repair.fill_holes(tm)
+    trimesh.repair.fix_normals(tm)
+    if verbose:
+        print(f"After hole filling: {len(tm.vertices)} vertices, {len(tm.faces)} faces")
+
+    # --- Step 2: Simplification ---
+    if decimation_target < len(tm.faces):
+        if HAS_FAST_SIMPLIFICATION:
+            ratio = min(1.0, decimation_target / max(len(tm.faces), 1))
+            new_verts, new_faces = fast_simplification.simplify(
+                tm.vertices.astype(np.float64), tm.faces,
+                target_reduction=(1.0 - ratio),
+            )
+            tm = trimesh.Trimesh(vertices=new_verts, faces=new_faces, process=False)
+        elif hasattr(tm, 'simplify_quadric_decimation'):
+            tm = tm.simplify_quadric_decimation(decimation_target)
+        if verbose:
+            print(f"After simplification: {len(tm.vertices)} vertices, {len(tm.faces)} faces")
+
+    trimesh.repair.fill_holes(tm)
+    trimesh.repair.fix_normals(tm)
+
+    if use_tqdm:
+        pbar.update(1)
+
+    # --- Step 3: UV Unwrapping with xatlas ---
+    if use_tqdm:
+        pbar.set_description("UV unwrapping")
+    if verbose:
+        print("UV unwrapping with xatlas...")
+
+    vmapping, out_faces_np, out_uvs_np = xatlas.parametrize(
+        tm.vertices.astype(np.float32),
+        tm.faces.astype(np.uint32),
+    )
+    out_vertices_np = tm.vertices[vmapping].astype(np.float32)
+    out_normals_np = tm.vertex_normals[vmapping].astype(np.float32)
+
+    if verbose:
+        print(f"After UV: {out_vertices_np.shape[0]} vertices, {out_faces_np.shape[0]} faces")
+    if use_tqdm:
+        pbar.update(1)
+
+    # --- Step 4: Texture Baking (GPU-accelerated) ---
+    if use_tqdm:
+        pbar.set_description("Baking textures (MPS)" if device.type == 'mps' else "Baking textures")
+    if verbose:
+        print("Baking textures...")
+
+    out_vertices_t = torch.from_numpy(out_vertices_np)
+    out_faces_t = torch.from_numpy(out_faces_np.astype(np.int64))
+    out_uvs_t = torch.from_numpy(out_uvs_np)
+
+    # GPU rasterization in UV space
+    pos, mask = _rasterize_uv_gpu(out_vertices_t, out_faces_t, out_uvs_t, texture_size, device)
+
+    valid_pos = torch.from_numpy(pos[mask]).float()
+
+    # Sample attributes from volume (GPU-accelerated)
+    attr_volume = attr_volume.cpu()
+    coords = coords.cpu()
+    attrs_full = torch.zeros(texture_size, texture_size, attr_volume.shape[1])
+    if valid_pos.shape[0] > 0:
+        sampled = _grid_sample_3d_gpu(attr_volume, coords, grid_size, valid_pos, voxel_size, aabb, device)
+        attrs_full[mask] = sampled
+
+    if use_tqdm:
+        pbar.update(1)
+
+    # --- Step 5: Texture Post-Processing ---
+    if use_tqdm:
+        pbar.set_description("Post-processing textures")
+    if verbose:
+        print("Post-processing textures...")
+
+    mask_inv = (~mask).astype(np.uint8)
+
+    base_color = np.clip(attrs_full[..., attr_layout['base_color']].numpy() * 255, 0, 255).astype(np.uint8)
+    metallic = np.clip(attrs_full[..., attr_layout['metallic']].numpy() * 255, 0, 255).astype(np.uint8)
+    roughness = np.clip(attrs_full[..., attr_layout['roughness']].numpy() * 255, 0, 255).astype(np.uint8)
+    alpha = np.clip(attrs_full[..., attr_layout['alpha']].numpy() * 255, 0, 255).astype(np.uint8)
+
+    # Auto-detect transparency from baked alpha values
+    alpha_valid = alpha[mask]
+    if alpha_valid.size > 0 and alpha_valid.min() < 250:
+        alpha_mode = 'BLEND'
+        if verbose:
+            print(f"Detected transparency (alpha min={alpha_valid.min()}), using BLEND mode")
+    else:
+        alpha_mode = 'OPAQUE'
+
+    # Inpainting to fill UV seams
+    base_color = cv2.inpaint(base_color, mask_inv, 3, cv2.INPAINT_TELEA)
+    metallic = cv2.inpaint(metallic, mask_inv, 1, cv2.INPAINT_TELEA)[..., None]
+    roughness = cv2.inpaint(roughness, mask_inv, 1, cv2.INPAINT_TELEA)[..., None]
+    alpha = cv2.inpaint(alpha, mask_inv, 1, cv2.INPAINT_TELEA)[..., None]
+
+    if use_tqdm:
+        pbar.update(1)
+
+    # --- Step 6: Build PBR Material & Export ---
+    if use_tqdm:
+        pbar.set_description("Finalizing")
+    if verbose:
+        print("Building PBR material...")
+
+    material = trimesh.visual.material.PBRMaterial(
+        baseColorTexture=Image.fromarray(np.concatenate([base_color, alpha], axis=-1)),
+        baseColorFactor=np.array([255, 255, 255, 255], dtype=np.uint8),
+        metallicRoughnessTexture=Image.fromarray(
+            np.concatenate([np.zeros_like(metallic), roughness, metallic], axis=-1)
+        ),
+        metallicFactor=1.0,
+        roughnessFactor=1.0,
+        alphaMode=alpha_mode,
+        doubleSided=True if not remesh else False,
+    )
+
+    # Coordinate system conversion (Y-up to Z-up for GLB)
+    vertices_out = out_vertices_np.copy()
+    normals_out = out_normals_np.copy()
+    vertices_out[:, 1], vertices_out[:, 2] = out_vertices_np[:, 2].copy(), -out_vertices_np[:, 1].copy()
+    normals_out[:, 1], normals_out[:, 2] = out_normals_np[:, 2].copy(), -out_normals_np[:, 1].copy()
+    uvs_out = out_uvs_np.copy()
+    uvs_out[:, 1] = 1 - uvs_out[:, 1]
+
+    textured_mesh = trimesh.Trimesh(
+        vertices=vertices_out,
+        faces=out_faces_np,
+        vertex_normals=normals_out,
+        process=False,
+        visual=trimesh.visual.TextureVisuals(uv=uvs_out, material=material),
+    )
+
+    if use_tqdm:
+        pbar.update(1)
+        pbar.close()
+    if verbose:
+        print("Done")
+
+    return textured_mesh

o-voxel/pyproject.toml

@@ -29,6 +29,16 @@ dependencies = [
     "tqdm",
     "zstandard",
     "easydict",
-    "cumesh @ git+https://github.com/JeffreyXiang/CuMesh.git",
-    "flex_gemm @ git+https://github.com/JeffreyXiang/FlexGEMM.git",
+]
+
+[project.optional-dependencies]
+gpu = [
+    "cumesh",
+    "flex_gemm",
+]
+metal = [
+    "mtldiffrast",
+]
+cuda = [
+    "nvdiffrast",
 ]

o-voxel/setup.py

@@ -1,27 +1,90 @@
 from setuptools import setup
-from torch.utils.cpp_extension import CUDAExtension, BuildExtension, IS_HIP_EXTENSION
 import os
+import platform
+import sys
 
 ROOT = os.path.dirname(os.path.abspath(__file__))
 BUILD_TARGET = os.environ.get("BUILD_TARGET", "auto")
 
-if BUILD_TARGET == "auto":
-    if IS_HIP_EXTENSION:
-        IS_HIP = True
+ext_modules = []
+cmdclass = {}
+
+if BUILD_TARGET == "cpu" or (BUILD_TARGET == "auto" and platform.system() == "Darwin"):
+    # CPU-only build: compile only C++ sources (Eigen-based), skip CUDA
+    from torch.utils.cpp_extension import CppExtension, BuildExtension
+
+    cpu_sources = [
+        "src/convert/flexible_dual_grid.cpp",
+        "src/convert/volumetic_attr.cpp",
+        "src/io/svo.cpp",
+        "src/io/filter_parent.cpp",
+        "src/io/filter_neighbor.cpp",
+        "src/ext_cpu.cpp",
+    ]
+    # Only build if the CPU ext entry point exists
+    ext_cpu_path = os.path.join(ROOT, "src/ext_cpu.cpp")
+    if os.path.exists(ext_cpu_path):
+        ext_modules = [
+            CppExtension(
+                name="o_voxel._C",
+                sources=cpu_sources,
+                include_dirs=[
+                    os.path.join(ROOT, "third_party/eigen"),
+                ],
+                extra_compile_args={
+                    "cxx": ["-O3", "-std=c++17"],
+                },
+            )
+        ]
     else:
-        IS_HIP = False
+        # No C++ extension — pure Python fallback
+        print("[o_voxel] Building without C++ extension (pure Python fallback)")
+        ext_modules = []
+    cmdclass = {'build_ext': BuildExtension}
 else:
-    if BUILD_TARGET == "cuda":
+    # CUDA/ROCm build (original)
+    from torch.utils.cpp_extension import CUDAExtension, BuildExtension, IS_HIP_EXTENSION
+
+    if BUILD_TARGET == "auto":
+        IS_HIP = IS_HIP_EXTENSION
+    elif BUILD_TARGET == "cuda":
         IS_HIP = False
     elif BUILD_TARGET == "rocm":
         IS_HIP = True
 
-if not IS_HIP:
+    if not IS_HIP:
         cc_flag = []
-else:
+    else:
         archs = os.getenv("GPU_ARCHS", "native").split(";")
         cc_flag = [f"--offload-arch={arch}" for arch in archs]
 
+    ext_modules = [
+        CUDAExtension(
+            name="o_voxel._C",
+            sources=[
+                "src/hash/hash.cu",
+                "src/convert/flexible_dual_grid.cpp",
+                "src/convert/volumetic_attr.cpp",
+                "src/serialize/api.cu",
+                "src/serialize/hilbert.cu",
+                "src/serialize/z_order.cu",
+                "src/io/svo.cpp",
+                "src/io/filter_parent.cpp",
+                "src/io/filter_neighbor.cpp",
+                "src/rasterize/rasterize.cu",
+                "src/ext.cpp",
+            ],
+            include_dirs=[
+                os.path.join(ROOT, "third_party/eigen"),
+            ],
+            extra_compile_args={
+                "cxx": ["-O3", "-std=c++17"],
+                "nvcc": ["-O3", "-std=c++17"] + cc_flag,
+            },
+        )
+    ]
+    cmdclass = {'build_ext': BuildExtension}
+
 setup(
     name="o_voxel",
     packages=[
@@ -29,39 +92,6 @@ setup(
         'o_voxel.convert',
         'o_voxel.io',
     ],
-    ext_modules=[
-        CUDAExtension(
-            name="o_voxel._C",
-            sources=[
-                # Hashmap functions
-                "src/hash/hash.cu",
-                # Convert functions
-                "src/convert/flexible_dual_grid.cpp",
-                "src/convert/volumetic_attr.cpp",
-                ## Serialization functions
-                "src/serialize/api.cu",
-                "src/serialize/hilbert.cu",
-                "src/serialize/z_order.cu",
-                # IO functions
-                "src/io/svo.cpp",
-                "src/io/filter_parent.cpp",
-                "src/io/filter_neighbor.cpp",
-                # Rasterization functions
-                "src/rasterize/rasterize.cu",
-                
-                # main
-                "src/ext.cpp",
-            ],
-            include_dirs=[
-                os.path.join(ROOT, "third_party/eigen"),
-            ],
-            extra_compile_args={
-                "cxx": ["-O3", "-std=c++17"],
-                "nvcc": ["-O3","-std=c++17"] + cc_flag,
-            }
-        )
-    ],
-    cmdclass={
-        'build_ext': BuildExtension
-    }
+    ext_modules=ext_modules,
+    cmdclass=cmdclass,
 )

requirements_macos.txt

@@ -0,0 +1,40 @@
+# Trellis2 MLX requirements (Apple Silicon)
+# NO: flash_attn, spconv, torchsparse (CUDA-only)
+
+# Core
+torch>=2.2.0
+torchvision>=0.17.0
+transformers>=4.40.0,<5
+safetensors
+huggingface_hub
+pillow
+numpy
+
+# Mesh processing
+trimesh
+xatlas
+fast-simplification
+pygltflib
+
+# Image processing
+opencv-python-headless
+
+# Server
+fastapi
+uvicorn[standard]
+pydantic>=2.0
+
+# MLX (Apple Silicon native)
+mlx>=0.31.0
+
+# Metal GPU packages (Apple Silicon — install with: pip install --no-build-isolation <url>)
+mtldiffrast @ https://github.com/pedronaugusto/mtldiffrast/archive/main.tar.gz
+cumesh @ https://github.com/pedronaugusto/mtlmesh/archive/main.tar.gz
+flex_gemm @ https://github.com/pedronaugusto/mtlgemm/archive/main.tar.gz
+
+# Misc
+tqdm
+imageio
+imageio-ffmpeg
+scipy
+einops

scripts/download_weights.py

@@ -0,0 +1,59 @@
+"""Download TRELLIS.2 weights into a local directory structure.
+
+The rest of the codebase already knows how to load TRELLIS.2 from a local
+snapshot directory as long as it contains the same layout as the Hugging Face
+repo root:
+
+    pipeline.json
+    texturing_pipeline.json
+    ckpts/*.json
+    ckpts/*.safetensors
+
+Example:
+    python export/download_weights.py --output-dir weights/TRELLIS.2-4B
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+
+from huggingface_hub import snapshot_download
+
+
+DEFAULT_PATTERNS = [
+    "pipeline.json",
+    "texturing_pipeline.json",
+    "ckpts/*",
+]
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Download TRELLIS.2 weights locally")
+    parser.add_argument("--repo-id", default="microsoft/TRELLIS.2-4B")
+    parser.add_argument("--output-dir", default="weights/TRELLIS.2-4B")
+    parser.add_argument(
+        "--full-snapshot",
+        action="store_true",
+        help="Download the entire repo instead of just pipeline configs and checkpoints",
+    )
+    args = parser.parse_args()
+
+    output_dir = os.path.abspath(args.output_dir)
+    os.makedirs(output_dir, exist_ok=True)
+
+    download_kwargs = {
+        "repo_id": args.repo_id,
+        "local_dir": output_dir,
+        "local_dir_use_symlinks": False,
+        "resume_download": True,
+    }
+    if not args.full_snapshot:
+        download_kwargs["allow_patterns"] = DEFAULT_PATTERNS
+
+    snapshot_path = snapshot_download(**download_kwargs)
+    print(f"Downloaded {args.repo_id} to {snapshot_path}")
+
+
+if __name__ == "__main__":
+    main()

scripts/start_server.sh

@@ -0,0 +1,23 @@
+#!/bin/bash
+# Start the Trellis2 MLX API server on macOS
+set -e
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+ROOT_DIR="$(dirname "$SCRIPT_DIR")"
+cd "$ROOT_DIR"
+
+# Set macOS-specific environment
+export SPARSE_CONV_BACKEND=pytorch
+export ATTN_BACKEND=sdpa
+export PYTORCH_ENABLE_MPS_FALLBACK=1  # deform_conv2d in RMBG not yet on MPS
+export PYTHONPATH="$ROOT_DIR/o-voxel:${PYTHONPATH:-}"
+
+# Defaults
+WEIGHTS="${TRELLIS2_WEIGHTS:-weights/TRELLIS.2-4B}"
+PORT="${TRELLIS2_PORT:-8082}"
+
+echo "Starting Trellis2 MLX API server..."
+echo "  Weights: $WEIGHTS"
+echo "  Port: $PORT"
+
+python api_server.py --weights "$WEIGHTS" --port "$PORT"

test_rast_parity.py

@@ -0,0 +1,148 @@
+"""
+Rasterization parity test for Trellis2 texture baking.
+Tests the exact UV-space rasterization scenario used by postprocess.py:
+  - Positions are UV coords mapped to clip space (z=0, w=1)
+  - All triangles share the same depth
+  - Each face must get its correct unique ID
+
+This is a focused unit test for the Metal rasterizer in the Trellis2 context.
+For full E2E parity (Metal vs CPU postprocess), see test_postprocess_parity.py.
+"""
+import torch
+import numpy as np
+import struct
+import pytest
+
+
+def float_to_triidx(f):
+    if f <= 16777216.0:
+        return int(f)
+    bits = struct.unpack('I', struct.pack('f', f))[0]
+    return bits - 0x4a800000
+
+
+@pytest.fixture
+def mtl_rast():
+    try:
+        from mtldiffrast.torch.ops import MtlRasterizeContext, rasterize
+        ctx = MtlRasterizeContext()
+        return ctx, rasterize
+    except ImportError:
+        pytest.skip("mtldiffrast not built")
+
+
+class TestTrellis2RastParity:
+    """UV-space rasterization as used by Trellis2 postprocess.py."""
+
+    def test_uv_quad_mesh(self, mtl_rast):
+        """Simple UV quad — the minimal texture baking scenario."""
+        ctx, rasterize = mtl_rast
+
+        # UV coords mapped to clip space: uv * 2 - 1, z=0, w=1
+        pos = torch.tensor([
+            [-1.0, -1.0, 0.0, 1.0],
+            [ 1.0, -1.0, 0.0, 1.0],
+            [ 1.0,  1.0, 0.0, 1.0],
+            [-1.0,  1.0, 0.0, 1.0],
+        ], dtype=torch.float32)
+        tri = torch.tensor([[0, 1, 2], [0, 2, 3]], dtype=torch.int32)
+
+        rast_out, _ = rasterize(ctx, pos, tri, resolution=[512, 512])
+        rast_np = rast_out[0].numpy()
+
+        # Full coverage
+        covered = (rast_np[:, :, 3] != 0).sum()
+        assert covered == 512 * 512, f"Expected full coverage, got {covered}"
+
+        # Both face IDs present
+        face_ids = set()
+        for py in range(0, 512, 32):
+            for px in range(0, 512, 32):
+                if rast_np[py, px, 3] != 0:
+                    fid = float_to_triidx(rast_np[py, px, 3]) - 1
+                    face_ids.add(fid)
+        assert face_ids == {0, 1}, f"Expected face IDs {{0, 1}}, got {face_ids}"
+
+    def test_random_uv_mesh(self, mtl_rast):
+        """Random UV layout with many triangles — simulates real Trellis2 mesh."""
+        ctx, rasterize = mtl_rast
+        torch.manual_seed(42)
+
+        # Generate a grid mesh with random UV perturbation
+        N = 20  # 20x20 grid = 800 triangles
+        verts = []
+        for y in range(N + 1):
+            for x in range(N + 1):
+                u = x / N
+                v = y / N
+                # Small random perturbation to avoid perfectly uniform spacing
+                if 0 < x < N and 0 < y < N:
+                    u += (torch.rand(1).item() - 0.5) * 0.02
+                    v += (torch.rand(1).item() - 0.5) * 0.02
+                cx = u * 2.0 - 1.0
+                cy = v * 2.0 - 1.0
+                verts.append([cx, cy, 0.0, 1.0])
+
+        tris = []
+        for y in range(N):
+            for x in range(N):
+                i = y * (N + 1) + x
+                tris.append([i, i + 1, i + N + 2])
+                tris.append([i, i + N + 2, i + N + 1])
+
+        pos = torch.tensor(verts, dtype=torch.float32)
+        tri = torch.tensor(tris, dtype=torch.int32)
+
+        rast_out, _ = rasterize(ctx, pos, tri, resolution=[256, 256])
+        rast_np = rast_out[0].numpy()
+
+        # Should have high coverage
+        covered = (rast_np[:, :, 3] != 0).sum()
+        total = 256 * 256
+        coverage_pct = covered / total * 100
+        assert coverage_pct > 90, f"Expected >90% coverage, got {coverage_pct:.1f}%"
+
+        # Should have many distinct face IDs (800 total)
+        face_ids = set()
+        for py in range(256):
+            for px in range(256):
+                if rast_np[py, px, 3] != 0:
+                    fid = float_to_triidx(rast_np[py, px, 3]) - 1
+                    face_ids.add(fid)
+        # At 256x256 resolution, most of the 800 triangles should appear
+        assert len(face_ids) > 600, f"Expected >600 distinct face IDs, got {len(face_ids)}"
+
+    def test_barycentrics_interpolation_valid(self, mtl_rast):
+        """Verify interpolated barycentrics are valid for texture sampling."""
+        ctx, rasterize = mtl_rast
+        from mtldiffrast.torch.ops import interpolate
+
+        pos = torch.tensor([
+            [-1.0, -1.0, 0.0, 1.0],
+            [ 1.0, -1.0, 0.0, 1.0],
+            [ 1.0,  1.0, 0.0, 1.0],
+            [-1.0,  1.0, 0.0, 1.0],
+        ], dtype=torch.float32)
+        tri = torch.tensor([[0, 1, 2], [0, 2, 3]], dtype=torch.int32)
+        # UV coordinates as vertex attributes
+        uv_attr = torch.tensor([
+            [0.0, 0.0, 0.0],
+            [1.0, 0.0, 0.0],
+            [1.0, 1.0, 0.0],
+            [0.0, 1.0, 0.0],
+        ], dtype=torch.float32)
+
+        rast_out, _ = rasterize(ctx, pos, tri, resolution=[64, 64])
+        interp_out, _ = interpolate(uv_attr, rast_out, tri)
+        interp_np = interp_out[0].numpy()
+
+        # Interpolated UVs should be in [0, 1]
+        mask = rast_out[0, :, :, 3].numpy() != 0
+        u_vals = interp_np[mask, 0]
+        v_vals = interp_np[mask, 1]
+        assert (u_vals >= -0.01).all() and (u_vals <= 1.01).all(), f"u out of range: [{u_vals.min()}, {u_vals.max()}]"
+        assert (v_vals >= -0.01).all() and (v_vals <= 1.01).all(), f"v out of range: [{v_vals.min()}, {v_vals.max()}]"
+
+
+if __name__ == '__main__':
+    pytest.main([__file__, '-v'])

trellis2/backends.py

@@ -0,0 +1,107 @@
+"""
+Centralized backend resolution for trellis2-apple.
+
+Resolves differentiable rasterization (dr), mesh processing (MeshBackend, BVH),
+remeshing, and grid_sample once at import time. All other modules import from here.
+
+Priority: Metal (mtldiffrast) > CUDA (nvdiffrast) > CPU fallback.
+Mesh: cumesh (auto-selects Metal/CUDA internally).
+"""
+import platform
+import torch
+
+# ---------------------------------------------------------------------------
+# Detect platform
+# ---------------------------------------------------------------------------
+HAS_MPS = hasattr(torch.backends, 'mps') and torch.backends.mps.is_available()
+HAS_CUDA = torch.cuda.is_available()
+
+# ---------------------------------------------------------------------------
+# Differentiable rasterizer  (dr)
+# mtldiffrast (Metal) and nvdiffrast (CUDA) are separate packages
+# ---------------------------------------------------------------------------
+dr = None
+_dr_backend = None
+
+try:
+    import mtldiffrast.torch as _mtldr
+    dr = _mtldr
+    _dr_backend = 'metal'
+except ImportError:
+    pass
+
+if dr is None:
+    try:
+        import nvdiffrast.torch as _nvdr
+        dr = _nvdr
+        _dr_backend = 'cuda'
+    except ImportError:
+        pass
+
+HAS_DR = dr is not None
+
+
+def RasterizeContext(device=None):
+    """Create the appropriate rasterization context for the active backend."""
+    if _dr_backend == 'metal':
+        return dr.MtlRasterizeContext(device=device)
+    elif _dr_backend == 'cuda':
+        return dr.RasterizeCudaContext(device=device)
+    raise RuntimeError("No differentiable rasterization backend available")
+
+# ---------------------------------------------------------------------------
+# Mesh processing  (MeshBackend, BVH, remesh)
+# cumesh auto-selects Metal or CUDA backend internally
+# ---------------------------------------------------------------------------
+MeshBackend = None
+BVH = None
+remesh_narrow_band_dc = None
+_mesh_backend = None
+
+try:
+    import cumesh
+    MeshBackend = cumesh.CuMesh
+    BVH = cumesh.cuBVH
+    remesh_narrow_band_dc = cumesh.remeshing.remesh_narrow_band_dc
+    _mesh_backend = 'metal' if platform.system() == 'Darwin' else 'cuda'
+except ImportError:
+    pass
+
+HAS_MESH = MeshBackend is not None
+HAS_REMESH = remesh_narrow_band_dc is not None
+
+# ---------------------------------------------------------------------------
+# CPU fallback mesh libraries
+# ---------------------------------------------------------------------------
+try:
+    import trimesh
+    HAS_TRIMESH = True
+except ImportError:
+    trimesh = None
+    HAS_TRIMESH = False
+
+try:
+    import fast_simplification
+    HAS_FAST_SIMPLIFICATION = True
+except ImportError:
+    fast_simplification = None
+    HAS_FAST_SIMPLIFICATION = False
+
+# ---------------------------------------------------------------------------
+# Grid sample (flex_gemm on Metal/CUDA, F.grid_sample fallback)
+# ---------------------------------------------------------------------------
+_flex_grid_sample_3d = None
+HAS_FLEX_GEMM = False
+
+try:
+    from flex_gemm.ops.grid_sample import grid_sample_3d as _fgs3d
+    _flex_grid_sample_3d = _fgs3d
+    HAS_FLEX_GEMM = True
+except ImportError:
+    pass
+
+# ---------------------------------------------------------------------------
+# Overall backend name
+# ---------------------------------------------------------------------------
+BACKEND = _dr_backend or _mesh_backend or ('cpu' if HAS_TRIMESH else None)
+HAS_GPU = HAS_MPS or HAS_CUDA

trellis2/modules/image_feature_extractor.py

@@ -1,4 +1,7 @@
 from typing import *
+import warnings
+warnings.filterwarnings("ignore", message="Importing from timm.models.layers is deprecated")
+warnings.filterwarnings("ignore", message="Importing from timm.models.registry is deprecated")
 import torch
 import torch.nn.functional as F
 from torchvision import transforms
@@ -18,15 +21,21 @@ class DinoV2FeatureExtractor:
         self.transform = transforms.Compose([
             transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
         ])
+        self._device = 'cpu'
 
     def to(self, device):
+        self._device = device
         self.model.to(device)
 
     def cuda(self):
-        self.model.cuda()
+        self.to('cuda')
 
     def cpu(self):
-        self.model.cpu()
+        self.to('cpu')
+
+    @property
+    def device(self):
+        return self._device
 
     @torch.no_grad()
     def __call__(self, image: Union[torch.Tensor, List[Image.Image]]) -> torch.Tensor:
@@ -46,11 +55,11 @@ class DinoV2FeatureExtractor:
             image = [i.resize((518, 518), Image.LANCZOS) for i in image]
             image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image]
             image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image]
-            image = torch.stack(image).cuda()
+            image = torch.stack(image).to(self._device)
         else:
             raise ValueError(f"Unsupported type of image: {type(image)}")
 
-        image = self.transform(image).cuda()
+        image = self.transform(image).to(self._device)
         features = self.model(image, is_training=True)['x_prenorm']
         patchtokens = F.layer_norm(features, features.shape[-1:])
         return patchtokens
@@ -68,15 +77,21 @@ class DinoV3FeatureExtractor:
         self.transform = transforms.Compose([
             transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
         ])
+        self._device = 'cpu'
 
     def to(self, device):
+        self._device = device
         self.model.to(device)
 
     def cuda(self):
-        self.model.cuda()
+        self.to('cuda')
 
     def cpu(self):
-        self.model.cpu()
+        self.to('cpu')
+
+    @property
+    def device(self):
+        return self._device
 
     def extract_features(self, image: torch.Tensor) -> torch.Tensor:
         image = image.to(self.model.embeddings.patch_embeddings.weight.dtype)
@@ -109,10 +124,10 @@ class DinoV3FeatureExtractor:
             image = [i.resize((self.image_size, self.image_size), Image.LANCZOS) for i in image]
             image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image]
             image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image]
-            image = torch.stack(image).cuda()
+            image = torch.stack(image).to(self._device)
         else:
             raise ValueError(f"Unsupported type of image: {type(image)}")
 
-        image = self.transform(image).cuda()
+        image = self.transform(image).to(self._device)
         features = self.extract_features(image)
         return features

trellis2/modules/sparse/attention/full_attn.py

@@ -211,6 +211,52 @@ def sparse_scaled_dot_product_attention(*args, **kwargs):
             max_q_seqlen = max(q_seqlen)
             max_kv_seqlen = max(kv_seqlen)
         out = flash_attn_3.flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_q_seqlen, max_kv_seqlen)
+    elif config.ATTN == 'sdpa':
+        import torch.nn.functional as F_attn
+        if num_all_args == 1:
+            q, k, v = qkv.unbind(dim=1)
+        elif num_all_args == 2:
+            k, v = kv.unbind(dim=1)
+        # Pad variable-length sequences into dense batch [N, max_len, H, C]
+        N = len(q_seqlen)
+        max_q = max(q_seqlen)
+        max_kv = max(kv_seqlen)
+        H = q.shape[-2]
+        C_q = q.shape[-1]
+        C_v = v.shape[-1]
+        # Build dense tensors
+        q_dense = q.new_zeros(N, max_q, H, C_q)
+        k_dense = k.new_zeros(N, max_kv, H, C_q)
+        v_dense = v.new_zeros(N, max_kv, H, C_v)
+        # Build attention mask
+        attn_mask = torch.zeros(N, max_q, max_kv, dtype=torch.bool, device=device)
+        q_offset = 0
+        kv_offset = 0
+        for i in range(N):
+            ql = q_seqlen[i]
+            kvl = kv_seqlen[i]
+            q_dense[i, :ql] = q[q_offset:q_offset + ql]
+            k_dense[i, :kvl] = k[kv_offset:kv_offset + kvl]
+            v_dense[i, :kvl] = v[kv_offset:kv_offset + kvl]
+            attn_mask[i, :ql, :kvl] = True
+            q_offset += ql
+            kv_offset += kvl
+        # sdpa expects [N, H, L, C], mask broadcastable to [N, H, Lq, Lkv]
+        q_dense = q_dense.permute(0, 2, 1, 3)  # [N, H, Lq, C]
+        k_dense = k_dense.permute(0, 2, 1, 3)  # [N, H, Lkv, C]
+        v_dense = v_dense.permute(0, 2, 1, 3)  # [N, H, Lkv, C_v]
+        # Expand mask for heads: [N, 1, Lq, Lkv]
+        sdpa_mask = attn_mask.unsqueeze(1)
+        # Use float mask for MPS compatibility (bool masks not supported)
+        float_mask = torch.zeros_like(sdpa_mask, dtype=q_dense.dtype)
+        float_mask.masked_fill_(~sdpa_mask, float('-inf'))
+        out_dense = F_attn.scaled_dot_product_attention(q_dense, k_dense, v_dense, attn_mask=float_mask)
+        # out_dense: [N, H, Lq, C_v] -> unpad back to packed
+        out_dense = out_dense.permute(0, 2, 1, 3)  # [N, Lq, H, C_v]
+        out_parts = []
+        for i in range(N):
+            out_parts.append(out_dense[i, :q_seqlen[i]])
+        out = torch.cat(out_parts, dim=0)
     else:
         raise ValueError(f"Unknown attention module: {config.ATTN}")
     

trellis2/modules/sparse/attention/windowed_attn.py

@@ -1,5 +1,6 @@
 from typing import *
 import torch
+import torch.nn.functional as F
 import math
 from .. import SparseTensor
 from .. import config
@@ -11,6 +12,97 @@ __all__ = [
 ]
 
 
+def _sdpa_varlen_qkvpacked(qkv_feats: torch.Tensor, attn_func_args: dict) -> torch.Tensor:
+    """sdpa for variable-length qkv-packed input (self-attention within windows)."""
+    q, k, v = qkv_feats.unbind(dim=1)  # each [M, H, C]
+    cu_seqlens = attn_func_args['cu_seqlens']
+    seq_lens = attn_func_args['seq_lens']
+    N = len(seq_lens)
+    max_len = attn_func_args['max_seqlen'].item() if isinstance(attn_func_args['max_seqlen'], torch.Tensor) else attn_func_args['max_seqlen']
+    H, C = q.shape[-2], q.shape[-1]
+
+    # Pad into dense batch [N, max_len, H, C]
+    q_dense = q.new_zeros(N, max_len, H, C)
+    k_dense = k.new_zeros(N, max_len, H, C)
+    v_dense = v.new_zeros(N, max_len, H, C)
+    mask = torch.zeros(N, max_len, dtype=torch.bool, device=q.device)
+    for i in range(N):
+        sl = seq_lens[i].item() if isinstance(seq_lens[i], torch.Tensor) else seq_lens[i]
+        start = cu_seqlens[i].item()
+        q_dense[i, :sl] = q[start:start + sl]
+        k_dense[i, :sl] = k[start:start + sl]
+        v_dense[i, :sl] = v[start:start + sl]
+        mask[i, :sl] = True
+
+    # [N, H, L, C]
+    q_dense = q_dense.permute(0, 2, 1, 3)
+    k_dense = k_dense.permute(0, 2, 1, 3)
+    v_dense = v_dense.permute(0, 2, 1, 3)
+
+    # Build float mask for MPS compatibility
+    sdpa_mask = mask.unsqueeze(1).unsqueeze(2)  # [N, 1, 1, L]
+    float_mask = torch.zeros(N, 1, max_len, max_len, dtype=q_dense.dtype, device=q.device)
+    float_mask.masked_fill_(~(mask.unsqueeze(1).unsqueeze(2) & mask.unsqueeze(1).unsqueeze(3)), float('-inf'))
+
+    out_dense = F.scaled_dot_product_attention(q_dense, k_dense, v_dense, attn_mask=float_mask)
+    out_dense = out_dense.permute(0, 2, 1, 3)  # [N, L, H, C]
+
+    # Unpad
+    parts = []
+    for i in range(N):
+        sl = seq_lens[i].item() if isinstance(seq_lens[i], torch.Tensor) else seq_lens[i]
+        parts.append(out_dense[i, :sl])
+    return torch.cat(parts, dim=0)
+
+
+def _sdpa_varlen(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
+                 q_args: dict, kv_args: dict) -> torch.Tensor:
+    """sdpa for variable-length cross-attention within windows."""
+    q_cu = q_args['cu_seqlens']
+    kv_cu = kv_args['cu_seqlens']
+    q_seq_lens = q_args['seq_lens']
+    kv_seq_lens = kv_args['seq_lens']
+    N = len(q_seq_lens)
+    max_q = q_args['max_seqlen'].item() if isinstance(q_args['max_seqlen'], torch.Tensor) else q_args['max_seqlen']
+    max_kv = kv_args['max_seqlen'].item() if isinstance(kv_args['max_seqlen'], torch.Tensor) else kv_args['max_seqlen']
+    H, C_q = q.shape[-2], q.shape[-1]
+    C_v = v.shape[-1]
+
+    q_dense = q.new_zeros(N, max_q, H, C_q)
+    k_dense = k.new_zeros(N, max_kv, H, C_q)
+    v_dense = v.new_zeros(N, max_kv, H, C_v)
+    q_mask = torch.zeros(N, max_q, dtype=torch.bool, device=q.device)
+    kv_mask = torch.zeros(N, max_kv, dtype=torch.bool, device=q.device)
+    for i in range(N):
+        ql = q_seq_lens[i].item() if isinstance(q_seq_lens[i], torch.Tensor) else q_seq_lens[i]
+        kvl = kv_seq_lens[i].item() if isinstance(kv_seq_lens[i], torch.Tensor) else kv_seq_lens[i]
+        qs = q_cu[i].item()
+        kvs = kv_cu[i].item()
+        q_dense[i, :ql] = q[qs:qs + ql]
+        k_dense[i, :kvl] = k[kvs:kvs + kvl]
+        v_dense[i, :kvl] = v[kvs:kvs + kvl]
+        q_mask[i, :ql] = True
+        kv_mask[i, :kvl] = True
+
+    q_dense = q_dense.permute(0, 2, 1, 3)
+    k_dense = k_dense.permute(0, 2, 1, 3)
+    v_dense = v_dense.permute(0, 2, 1, 3)
+
+    # q_mask: (N, max_q), kv_mask: (N, max_kv) -> cross mask: (N, 1, max_q, max_kv)
+    cross_mask = q_mask.unsqueeze(2) & kv_mask.unsqueeze(1)  # (N, max_q, max_kv)
+    float_mask = torch.zeros(N, 1, max_q, max_kv, dtype=q_dense.dtype, device=q.device)
+    float_mask.masked_fill_(~cross_mask.unsqueeze(1), float('-inf'))
+
+    out_dense = F.scaled_dot_product_attention(q_dense, k_dense, v_dense, attn_mask=float_mask)
+    out_dense = out_dense.permute(0, 2, 1, 3)
+
+    parts = []
+    for i in range(N):
+        ql = q_seq_lens[i].item() if isinstance(q_seq_lens[i], torch.Tensor) else q_seq_lens[i]
+        parts.append(out_dense[i, :ql])
+    return torch.cat(parts, dim=0)
+
+
 def calc_window_partition(
     tensor: SparseTensor,
     window_size: Union[int, Tuple[int, ...]],
@@ -60,6 +152,12 @@ def calc_window_partition(
             'cu_seqlens': torch.cat([torch.tensor([0], device=tensor.device), torch.cumsum(seq_lens, dim=0)], dim=0).int(),
             'max_seqlen': torch.max(seq_lens)
         }
+    elif config.ATTN == 'sdpa':
+        attn_func_args = {
+            'cu_seqlens': torch.cat([torch.tensor([0], device=tensor.device), torch.cumsum(seq_lens, dim=0)], dim=0).int(),
+            'max_seqlen': torch.max(seq_lens),
+            'seq_lens': seq_lens,
+        }
 
     return fwd_indices, bwd_indices, seq_lens, attn_func_args
     
@@ -113,6 +211,8 @@ def sparse_windowed_scaled_dot_product_self_attention(
         if 'flash_attn' not in globals():
             import flash_attn
         out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv_feats, **attn_func_args)  # [M, H, C]
+    elif config.ATTN == 'sdpa':
+        out = _sdpa_varlen_qkvpacked(qkv_feats, attn_func_args)                         # [M, H, C]
 
     out = out[bwd_indices]      # [T, H, C]
 
@@ -172,11 +272,11 @@ def sparse_windowed_scaled_dot_product_cross_attention(
         if 'xops' not in globals():
             import xformers.ops as xops
         k, v = kv_feats.unbind(dim=1)                                                   # [M, H, C]
-        q = q.unsqueeze(0)                                                              # [1, M, H, C]
+        q_feats_u = q_feats.unsqueeze(0)                                                # [1, M, H, C]
         k = k.unsqueeze(0)                                                              # [1, M, H, C]
         v = v.unsqueeze(0)                                                              # [1, M, H, C]
         mask = xops.fmha.BlockDiagonalMask.from_seqlens(q_seq_lens, kv_seq_lens)
-        out = xops.memory_efficient_attention(q, k, v, attn_bias=mask)[0]               # [M, H, C]
+        out = xops.memory_efficient_attention(q_feats_u, k, v, attn_bias=mask)[0]       # [M, H, C]
     elif config.ATTN == 'flash_attn':
         if 'flash_attn' not in globals():
             import flash_attn
@@ -184,6 +284,9 @@ def sparse_windowed_scaled_dot_product_cross_attention(
             cu_seqlens_q=q_attn_func_args['cu_seqlens'], cu_seqlens_k=kv_attn_func_args['cu_seqlens'],
             max_seqlen_q=q_attn_func_args['max_seqlen'], max_seqlen_k=kv_attn_func_args['max_seqlen'],
         )  # [M, H, C]
+    elif config.ATTN == 'sdpa':
+        k, v = kv_feats.unbind(dim=1)
+        out = _sdpa_varlen(q_feats, k, v, q_attn_func_args, kv_attn_func_args)          # [M, H, C]
 
     out = out[q_bwd_indices]      # [T, H, C]
 

trellis2/modules/sparse/config.py

@@ -1,9 +1,32 @@
 from typing import *
+import platform
+import sys
 
 CONV = 'flex_gemm'
 DEBUG = False
 ATTN = 'flash_attn'
 
+def __detect_defaults():
+    """Auto-detect best backends for current platform."""
+    global CONV, ATTN
+    if platform.system() == 'Darwin':
+        ATTN = 'sdpa'
+        try:
+            import flex_gemm
+            CONV = 'flex_gemm'
+        except ImportError:
+            CONV = 'pytorch'
+    elif not __has_cuda():
+        CONV = 'pytorch'
+        ATTN = 'sdpa'
+
+def __has_cuda():
+    try:
+        import torch
+        return torch.cuda.is_available()
+    except Exception:
+        return False
+
 def __from_env():
     import os
 
@@ -11,17 +34,19 @@ def __from_env():
     global DEBUG
     global ATTN
 
+    __detect_defaults()
+
     env_sparse_conv_backend = os.environ.get('SPARSE_CONV_BACKEND')
     env_sparse_debug = os.environ.get('SPARSE_DEBUG')
     env_sparse_attn_backend = os.environ.get('SPARSE_ATTN_BACKEND')
     if env_sparse_attn_backend is None:
         env_sparse_attn_backend = os.environ.get('ATTN_BACKEND')
 
-    if env_sparse_conv_backend is not None and env_sparse_conv_backend in ['none', 'spconv', 'torchsparse', 'flex_gemm']:
+    if env_sparse_conv_backend is not None and env_sparse_conv_backend in ['none', 'spconv', 'torchsparse', 'flex_gemm', 'pytorch']:
         CONV = env_sparse_conv_backend
     if env_sparse_debug is not None:
         DEBUG = env_sparse_debug == '1'
-    if env_sparse_attn_backend is not None and env_sparse_attn_backend in ['xformers', 'flash_attn', 'flash_attn_3']:
+    if env_sparse_attn_backend is not None and env_sparse_attn_backend in ['xformers', 'flash_attn', 'flash_attn_3', 'sdpa']:
         ATTN = env_sparse_attn_backend
 
     print(f"[SPARSE] Conv backend: {CONV}; Attention backend: {ATTN}")
@@ -30,7 +55,7 @@ def __from_env():
 __from_env()
 
 
-def set_conv_backend(backend: Literal['none', 'spconv', 'torchsparse', 'flex_gemm']):
+def set_conv_backend(backend: Literal['none', 'spconv', 'torchsparse', 'flex_gemm', 'pytorch']):
     global CONV
     CONV = backend
 
@@ -38,6 +63,6 @@ def set_debug(debug: bool):
     global DEBUG
     DEBUG = debug
 
-def set_attn_backend(backend: Literal['xformers', 'flash_attn']):
+def set_attn_backend(backend: Literal['xformers', 'flash_attn', 'sdpa']):
     global ATTN
     ATTN = backend

trellis2/modules/sparse/conv/conv_pytorch.py

@@ -0,0 +1,174 @@
+"""
+Pure PyTorch submanifold sparse conv3d backend.
+
+Algorithm for 3×3×3 submanifold sparse conv:
+1. Hash all voxel coords → flat index (b×S³ + x×S² + y×S + z)
+2. For each of 27 kernel offsets: shift coords, lookup neighbors via hash
+3. Gather neighbor features, matmul with kernel weight per offset, scatter_add
+4. Cache neighbor maps (topology constant during inference)
+"""
+import math
+import itertools
+import torch
+import torch.nn as nn
+from .. import SparseTensor
+
+
+def _build_neighbor_map(coords: torch.Tensor, shape: torch.Size, spatial_shape: torch.Size,
+                        kernel_size: tuple, dilation: tuple, device: torch.device):
+    """
+    Build neighbor map for submanifold sparse conv.
+
+    Returns:
+        neighbor_map: (K, N) int64 tensor. For each kernel offset k and voxel i,
+                      neighbor_map[k, i] = index of neighbor in coords, or -1 if absent.
+    """
+    N = coords.shape[0]
+    batch_size = shape[0]
+    D, H, W = spatial_shape
+
+    # Ensure coords are on the target device
+    if coords.device != device:
+        coords = coords.to(device)
+
+    # Build lookup table: flat_coord -> index in coords tensor
+    # Flat key = batch * (D*H*W) + x*H*W + y*W + z
+    DHW = D * H * W
+    flat_keys = (coords[:, 0].long() * DHW +
+                 coords[:, 1].long() * (H * W) +
+                 coords[:, 2].long() * W +
+                 coords[:, 3].long())
+    table_size = batch_size * DHW
+    lookup = torch.full((table_size,), -1, dtype=torch.long, device=device)
+    lookup[flat_keys] = torch.arange(N, dtype=torch.long, device=device)
+
+    # Generate kernel offsets
+    kd, kh, kw = kernel_size
+    dd, dh, dw = dilation
+    offsets = []
+    for dx, dy, dz in itertools.product(
+        range(-(kd // 2), kd // 2 + 1),
+        range(-(kh // 2), kh // 2 + 1),
+        range(-(kw // 2), kw // 2 + 1),
+    ):
+        offsets.append((dx * dd, dy * dh, dz * dw))
+    K = len(offsets)
+    offsets_tensor = torch.tensor(offsets, dtype=coords.dtype, device=device)  # (K, 3)
+
+    # For each offset, shift coords and look up
+    neighbor_map = torch.full((K, N), -1, dtype=torch.long, device=device)
+    for k_idx in range(K):
+        shifted = coords.clone()
+        shifted[:, 1] = shifted[:, 1] + offsets_tensor[k_idx, 0]
+        shifted[:, 2] = shifted[:, 2] + offsets_tensor[k_idx, 1]
+        shifted[:, 3] = shifted[:, 3] + offsets_tensor[k_idx, 2]
+
+        # Bounds check
+        valid = ((shifted[:, 1] >= 0) & (shifted[:, 1] < D) &
+                 (shifted[:, 2] >= 0) & (shifted[:, 2] < H) &
+                 (shifted[:, 3] >= 0) & (shifted[:, 3] < W))
+
+        flat_shifted = (shifted[:, 0].long() * DHW +
+                        shifted[:, 1].long() * (H * W) +
+                        shifted[:, 2].long() * W +
+                        shifted[:, 3].long())
+        # Clamp for safe indexing (invalid entries will be masked out)
+        flat_shifted = flat_shifted.clamp(0, table_size - 1)
+        looked_up = lookup[flat_shifted]
+        looked_up[~valid] = -1
+        neighbor_map[k_idx] = looked_up
+
+    return neighbor_map
+
+
+def sparse_conv3d_init(self, in_channels, out_channels, kernel_size, stride=1,
+                       dilation=1, padding=None, bias=True, indice_key=None):
+    assert stride == 1 and (padding is None), \
+        'PyTorch backend only supports submanifold sparse convolution (stride=1, padding=None)'
+
+    self.in_channels = in_channels
+    self.out_channels = out_channels
+    self.kernel_size = tuple(kernel_size) if isinstance(kernel_size, (list, tuple)) else (kernel_size,) * 3
+    self.stride = tuple(stride) if isinstance(stride, (list, tuple)) else (stride,) * 3
+    self.dilation = tuple(dilation) if isinstance(dilation, (list, tuple)) else (dilation,) * 3
+
+    # Store weight in same format as flex_gemm: (Co, Kd, Kh, Kw, Ci)
+    # This ensures checkpoint compatibility — flex_gemm permutes (Co, Ci, K...) -> (Co, K..., Ci) at init
+    self.weight = nn.Parameter(torch.empty((out_channels, *self.kernel_size, in_channels)))
+    if bias:
+        self.bias = nn.Parameter(torch.empty(out_channels))
+    else:
+        self.register_parameter("bias", None)
+
+    # Initialize parameters
+    torch.nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+    if self.bias is not None:
+        fan_in, _ = torch.nn.init._calculate_fan_in_and_fan_out(
+            self.weight.permute(0, 4, 1, 2, 3))  # back to (Co, Ci, K...) for fan calc
+        if fan_in != 0:
+            bound = 1 / math.sqrt(fan_in)
+            torch.nn.init.uniform_(self.bias, -bound, bound)
+
+
+def _get_weight_kernel(self):
+    """Reshape stored weight (Co, Kd, Kh, Kw, Ci) -> (K, Ci, Co) for gather-matmul."""
+    Co, Kd, Kh, Kw, Ci = self.weight.shape
+    K = Kd * Kh * Kw
+    # (Co, Kd, Kh, Kw, Ci) -> (Co, K, Ci) -> (K, Ci, Co)
+    return self.weight.reshape(Co, K, Ci).permute(1, 2, 0)
+
+
+def sparse_conv3d_forward(self, x: SparseTensor) -> SparseTensor:
+    N = x.feats.shape[0]
+
+    # Check cache for neighbor map
+    Co, Kd, Kh, Kw, Ci = self.weight.shape
+    cache_key = f'SubMConv3d_neighbor_cache_pytorch_{self.kernel_size}_dilation{self.dilation}'
+    neighbor_map = x.get_spatial_cache(cache_key)
+
+    if neighbor_map is None:
+        neighbor_map = _build_neighbor_map(
+            x.coords, x.shape, x.spatial_shape,
+            self.kernel_size, self.dilation, x.device,
+        )
+        x.register_spatial_cache(cache_key, neighbor_map)
+
+    K = neighbor_map.shape[0]
+    w = _get_weight_kernel(self)  # (K, Ci, Co)
+
+    # Pad feats with a zero row for -1 indices
+    feats_padded = torch.cat([x.feats, torch.zeros(1, x.feats.shape[1], device=x.device, dtype=x.feats.dtype)], dim=0)
+    pad_idx = N
+
+    # Replace -1 with pad_idx
+    safe_map = neighbor_map.clone()
+    safe_map[safe_map < 0] = pad_idx  # (K, N)
+
+    # Gather: (K, N, Ci)
+    gathered = feats_padded[safe_map]
+
+    # Matmul per kernel offset: (K, N, Ci) @ (K, Ci, Co) -> (K, N, Co)
+    out = torch.bmm(gathered, w.to(gathered.dtype))
+
+    # Mask out invalid neighbors before summing
+    valid_mask = (neighbor_map >= 0).unsqueeze(-1)  # (K, N, 1)
+    out = out * valid_mask
+
+    # Sum over kernel offsets
+    result = out.sum(dim=0)  # (N, Co)
+
+    if self.bias is not None:
+        result = result + self.bias
+
+    return x.replace(result)
+
+
+def sparse_inverse_conv3d_init(self, in_channels, out_channels, kernel_size, stride=1,
+                                dilation=1, bias=True, indice_key=None):
+    sparse_conv3d_init(self, in_channels, out_channels, kernel_size, stride=1,
+                       dilation=dilation, padding=None, bias=bias, indice_key=indice_key)
+
+
+def sparse_inverse_conv3d_forward(self, x: SparseTensor) -> SparseTensor:
+    # For submanifold case (stride=1), inverse conv is the same as forward conv
+    return sparse_conv3d_forward(self, x)

trellis2/pipelines/rembg/BiRefNet.py

@@ -8,7 +8,7 @@ from PIL import Image
 class BiRefNet:
     def __init__(self, model_name: str = "ZhengPeng7/BiRefNet"):
         self.model = AutoModelForImageSegmentation.from_pretrained(
-            model_name, trust_remote_code=True
+            model_name, trust_remote_code=True,
         )
         self.model.eval()
         self.transform_image = transforms.Compose(
@@ -18,20 +18,21 @@ class BiRefNet:
                 transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
             ]
         )
+        self._device = 'cpu'
 
-    def to(self, device: str):
+    def to(self, device):
+        self._device = device
         self.model.to(device)
 
     def cuda(self):
-        self.model.cuda()
+        self.to('cuda')
 
     def cpu(self):
-        self.model.cpu()
+        self.to('cpu')
 
     def __call__(self, image: Image.Image) -> Image.Image:
         image_size = image.size
-        input_images = self.transform_image(image).unsqueeze(0).to("cuda")
-        # Prediction
+        input_images = self.transform_image(image).unsqueeze(0).to(self._device)
         with torch.no_grad():
             preds = self.model(input_images)[-1].sigmoid().cpu()
         pred = preds[0].squeeze()

trellis2/pipelines/trellis2_image_to_3d.py

@@ -1,4 +1,5 @@
 from typing import *
+import gc
 import torch
 import torch.nn as nn
 import numpy as np
@@ -10,6 +11,15 @@ from ..modules import image_feature_extractor
 from ..representations import Mesh, MeshWithVoxel
 
 
+def _free_memory():
+    """Aggressively free memory after model offload."""
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    elif hasattr(torch, 'mps') and hasattr(torch.mps, 'empty_cache'):
+        torch.mps.empty_cache()
+
+
 class Trellis2ImageTo3DPipeline(Pipeline):
     """
     Pipeline for inferring Trellis2 image-to-3D models.
@@ -147,6 +157,7 @@ class Trellis2ImageTo3DPipeline(Pipeline):
             output = self.rembg_model(input)
             if self.low_vram:
                 self.rembg_model.cpu()
+                _free_memory()
         output_np = np.array(output)
         alpha = output_np[:, :, 3]
         bbox = np.argwhere(alpha > 0.8 * 255)
@@ -177,6 +188,7 @@ class Trellis2ImageTo3DPipeline(Pipeline):
         cond = self.image_cond_model(image)
         if self.low_vram:
             self.image_cond_model.cpu()
+            _free_memory()
         if not include_neg_cond:
             return {'cond': cond}
         neg_cond = torch.zeros_like(cond)
@@ -219,6 +231,7 @@ class Trellis2ImageTo3DPipeline(Pipeline):
         ).samples
         if self.low_vram:
             flow_model.cpu()
+            _free_memory()
 
         # Decode sparse structure latent
         decoder = self.models['sparse_structure_decoder']
@@ -227,6 +240,7 @@ class Trellis2ImageTo3DPipeline(Pipeline):
         decoded = decoder(z_s)>0
         if self.low_vram:
             decoder.cpu()
+            _free_memory()
         if resolution != decoded.shape[2]:
             ratio = decoded.shape[2] // resolution
             decoded = torch.nn.functional.max_pool3d(decoded.float(), ratio, ratio, 0) > 0.5
@@ -267,6 +281,7 @@ class Trellis2ImageTo3DPipeline(Pipeline):
         ).samples
         if self.low_vram:
             flow_model.cpu()
+            _free_memory()
 
         std = torch.tensor(self.shape_slat_normalization['std'])[None].to(slat.device)
         mean = torch.tensor(self.shape_slat_normalization['mean'])[None].to(slat.device)
@@ -312,6 +327,7 @@ class Trellis2ImageTo3DPipeline(Pipeline):
         ).samples
         if self.low_vram:
             flow_model_lr.cpu()
+            _free_memory()
         std = torch.tensor(self.shape_slat_normalization['std'])[None].to(slat.device)
         mean = torch.tensor(self.shape_slat_normalization['mean'])[None].to(slat.device)
         slat = slat * std + mean
@@ -324,6 +340,7 @@ class Trellis2ImageTo3DPipeline(Pipeline):
         if self.low_vram:
             self.models['shape_slat_decoder'].cpu()
             self.models['shape_slat_decoder'].low_vram = False
+            _free_memory()
         hr_resolution = resolution
         while True:
             quant_coords = torch.cat([
@@ -356,6 +373,7 @@ class Trellis2ImageTo3DPipeline(Pipeline):
         ).samples
         if self.low_vram:
             flow_model.cpu()
+            _free_memory()
 
         std = torch.tensor(self.shape_slat_normalization['std'])[None].to(slat.device)
         mean = torch.tensor(self.shape_slat_normalization['mean'])[None].to(slat.device)
@@ -386,6 +404,7 @@ class Trellis2ImageTo3DPipeline(Pipeline):
         if self.low_vram:
             self.models['shape_slat_decoder'].cpu()
             self.models['shape_slat_decoder'].low_vram = False
+            _free_memory()
         return ret
 
     def sample_tex_slat(
@@ -424,6 +443,7 @@ class Trellis2ImageTo3DPipeline(Pipeline):
         ).samples
         if self.low_vram:
             flow_model.cpu()
+            _free_memory()
 
         std = torch.tensor(self.tex_slat_normalization['std'])[None].to(slat.device)
         mean = torch.tensor(self.tex_slat_normalization['mean'])[None].to(slat.device)
@@ -450,6 +470,7 @@ class Trellis2ImageTo3DPipeline(Pipeline):
         ret = self.models['tex_slat_decoder'](slat, guide_subs=subs) * 0.5 + 0.5
         if self.low_vram:
             self.models['tex_slat_decoder'].cpu()
+            _free_memory()
         return ret
     
     @torch.no_grad()
@@ -587,7 +608,10 @@ class Trellis2ImageTo3DPipeline(Pipeline):
                 cond_1024, self.models['tex_slat_flow_model_1024'],
                 shape_slat, tex_slat_sampler_params
             )
+        if torch.cuda.is_available():
             torch.cuda.empty_cache()
+        elif hasattr(torch, 'mps') and hasattr(torch.mps, 'empty_cache'):
+            torch.mps.empty_cache()
         out_mesh = self.decode_latent(shape_slat, tex_slat, res)
         if return_latent:
             return out_mesh, (shape_slat, tex_slat, res)

trellis2/pipelines/trellis2_texturing.py

@@ -9,10 +9,12 @@ from . import samplers, rembg
 from ..modules.sparse import SparseTensor
 from ..modules import image_feature_extractor
 import o_voxel
-import cumesh
-import nvdiffrast.torch as dr
 import cv2
-import flex_gemm
+from ..utils.grid_sample import grid_sample_3d as _grid_sample_3d
+
+from ..backends import dr, RasterizeContext, MeshBackend as _MeshBackend
+if dr is None or _MeshBackend is None:
+    raise ImportError("cumesh + mtldiffrast (macOS) or cumesh + nvdiffrast (CUDA) required")
 
 
 class Trellis2TexturingPipeline(Pipeline):
@@ -294,26 +296,26 @@ class Trellis2TexturingPipeline(Pipeline):
         vertices = mesh.vertices
         faces = mesh.faces
         normals = mesh.vertex_normals
-        vertices_torch = torch.from_numpy(vertices).float().cuda()
-        faces_torch = torch.from_numpy(faces).int().cuda()
+        vertices_torch = torch.from_numpy(vertices).float().to(self.device)
+        faces_torch = torch.from_numpy(faces).int().to(self.device)
         if hasattr(mesh, 'visual') and hasattr(mesh.visual, 'uv') and mesh.visual.uv is not None:
             uvs = mesh.visual.uv.copy()
             uvs[:, 1] = 1 - uvs[:, 1]
-            uvs_torch = torch.from_numpy(uvs).float().cuda()
+            uvs_torch = torch.from_numpy(uvs).float().to(self.device)
         else:
-            _cumesh = cumesh.CuMesh()
-            _cumesh.init(vertices_torch, faces_torch)
-            vertices_torch, faces_torch, uvs_torch, vmap = _cumesh.uv_unwrap(return_vmaps=True)
-            vertices_torch = vertices_torch.cuda()
-            faces_torch = faces_torch.cuda()
-            uvs_torch = uvs_torch.cuda()
+            _mesh_op = _MeshBackend()
+            _mesh_op.init(vertices_torch, faces_torch)
+            vertices_torch, faces_torch, uvs_torch, vmap = _mesh_op.uv_unwrap(return_vmaps=True)
+            vertices_torch = vertices_torch.to(self.device)
+            faces_torch = faces_torch.to(self.device)
+            uvs_torch = uvs_torch.to(self.device)
             vertices = vertices_torch.cpu().numpy()
             faces = faces_torch.cpu().numpy()
             uvs = uvs_torch.cpu().numpy()
             normals = normals[vmap.cpu().numpy()]
                 
         # rasterize
-        ctx = dr.RasterizeCudaContext()
+        ctx = RasterizeContext()
         uvs_torch = torch.cat([uvs_torch * 2 - 1, torch.zeros_like(uvs_torch[:, :1]), torch.ones_like(uvs_torch[:, :1])], dim=-1).unsqueeze(0)
         rast, _ = dr.rasterize(
             ctx, uvs_torch, faces_torch,
@@ -323,7 +325,7 @@ class Trellis2TexturingPipeline(Pipeline):
         pos = dr.interpolate(vertices_torch.unsqueeze(0), rast, faces_torch)[0][0]
         
         attrs = torch.zeros(texture_size, texture_size, pbr_voxel.shape[1], device=self.device)
-        attrs[mask] = flex_gemm.ops.grid_sample.grid_sample_3d(
+        attrs[mask] = _grid_sample_3d(
             pbr_voxel.feats,
             pbr_voxel.coords,
             shape=torch.Size([*pbr_voxel.shape, *pbr_voxel.spatial_shape]),

trellis2/renderers/mesh_renderer.py

@@ -3,6 +3,7 @@ import torch
 from easydict import EasyDict as edict
 from ..representations.mesh import Mesh, MeshWithVoxel, MeshWithPbrMaterial, TextureFilterMode, AlphaMode, TextureWrapMode
 import torch.nn.functional as F
+from ..backends import dr, RasterizeContext
 
 
 def intrinsics_to_projection(
@@ -41,9 +42,6 @@ class MeshRenderer:
         rendering_options (dict): Rendering options.
         """
     def __init__(self, rendering_options={}, device='cuda'):
-        if 'dr' not in globals():
-            import nvdiffrast.torch as dr
-        
         self.rendering_options = edict({
             "resolution": None,
             "near": None,
@@ -54,7 +52,7 @@ class MeshRenderer:
             "clamp_barycentric_coords": False,
         })
         self.rendering_options.update(rendering_options)
-        self.glctx = dr.RasterizeCudaContext(device=device)
+        self.glctx = RasterizeContext(device=device)
         self.device=device
         
     def render(
@@ -81,9 +79,6 @@ class MeshRenderer:
                 normal (torch.Tensor): [3, H, W] rendered normal image
                 mask (torch.Tensor): [H, W] rendered mask image
         """
-        if 'dr' not in globals():
-            import nvdiffrast.torch as dr
-            
         resolution = self.rendering_options["resolution"]
         near = self.rendering_options["near"]
         far = self.rendering_options["far"]
@@ -159,12 +154,11 @@ class MeshRenderer:
                     if antialias: img = dr.antialias(img, rast, vertices_clip, faces)
                 elif type == "attr":
                     if isinstance(mesh, MeshWithVoxel):
-                        if 'grid_sample_3d' not in globals():
-                            from flex_gemm.ops.grid_sample import grid_sample_3d
+                        from ..utils.grid_sample import grid_sample_3d as _grid_sample_3d
                         mask = rast[..., -1:] > 0
                         xyz = dr.interpolate(vertices, rast, faces)[0]
                         xyz = ((xyz - mesh.origin) / mesh.voxel_size).reshape(1, -1, 3)
-                        img = grid_sample_3d(
+                        img = _grid_sample_3d(
                             mesh.attrs,
                             torch.cat([torch.zeros_like(mesh.coords[..., :1]), mesh.coords], dim=-1),
                             mesh.voxel_shape,
@@ -290,12 +284,11 @@ class MeshRenderer:
                         img = dr.interpolate(vertices, rast, faces_chunk)[0]
                     elif type == "attr":
                         if isinstance(mesh, MeshWithVoxel):
-                            if 'grid_sample_3d' not in globals():
-                                from flex_gemm.ops.grid_sample import grid_sample_3d
+                            from ..utils.grid_sample import grid_sample_3d as _grid_sample_3d
                             mask = rast[..., -1:] > 0
                             xyz = dr.interpolate(vertices, rast, faces_chunk)[0]
                             xyz = ((xyz - mesh.origin) / mesh.voxel_size).reshape(1, -1, 3)
-                            img = grid_sample_3d(
+                            img = _grid_sample_3d(
                                 mesh.attrs,
                                 torch.cat([torch.zeros_like(mesh.coords[..., :1]), mesh.coords], dim=-1),
                                 mesh.voxel_shape,

trellis2/renderers/pbr_mesh_renderer.py

@@ -5,6 +5,7 @@ import numpy as np
 import utils3d
 from ..representations.mesh import Mesh, MeshWithVoxel, MeshWithPbrMaterial, TextureFilterMode, AlphaMode, TextureWrapMode
 import torch.nn.functional as F
+from ..backends import dr, RasterizeContext
 
 
 def cube_to_dir(s, x, y):
@@ -18,12 +19,11 @@ def cube_to_dir(s, x, y):
 
 
 def latlong_to_cubemap(latlong_map, res):
-    if 'dr' not in globals():
-        import nvdiffrast.torch as dr
-    cubemap = torch.zeros(6, res[0], res[1], latlong_map.shape[-1], dtype=torch.float32, device='cuda')
+    device = latlong_map.device
+    cubemap = torch.zeros(6, res[0], res[1], latlong_map.shape[-1], dtype=torch.float32, device=device)
     for s in range(6):
-        gy, gx = torch.meshgrid(torch.linspace(-1.0 + 1.0 / res[0], 1.0 - 1.0 / res[0], res[0], device='cuda'), 
-                                torch.linspace(-1.0 + 1.0 / res[1], 1.0 - 1.0 / res[1], res[1], device='cuda'),
+        gy, gx = torch.meshgrid(torch.linspace(-1.0 + 1.0 / res[0], 1.0 - 1.0 / res[0], res[0], device=device),
+                                torch.linspace(-1.0 + 1.0 / res[1], 1.0 - 1.0 / res[1], res[1], device=device),
                                 indexing='ij')
         v = F.normalize(cube_to_dir(s, gx, gy), dim=-1)
 
@@ -53,8 +53,6 @@ class EnvMap:
         return self._backend.shade(gb_pos, gb_normal, kd, ks, view_pos, specular)
     
     def sample(self, directions: torch.Tensor):
-        if 'dr' not in globals():
-            import nvdiffrast.torch as dr
         return dr.texture(
             self._backend.base.unsqueeze(0),
             directions.unsqueeze(0),
@@ -195,9 +193,6 @@ class PbrMeshRenderer:
         rendering_options (dict): Rendering options.
         """
     def __init__(self, rendering_options={}, device='cuda'):
-        if 'dr' not in globals():
-            import nvdiffrast.torch as dr
-        
         self.rendering_options = edict({
             "resolution": None,
             "near": None,
@@ -206,7 +201,7 @@ class PbrMeshRenderer:
             "peel_layers": 8,
         })
         self.rendering_options.update(rendering_options)
-        self.glctx = dr.RasterizeCudaContext(device=device)
+        self.glctx = RasterizeContext(device=device)
         self.device=device
         
     def render(
@@ -237,9 +232,6 @@ class PbrMeshRenderer:
                 metallic (torch.Tensor): [H, W] metallic image
                 roughness (torch.Tensor): [H, W] roughness image
         """
-        if 'dr' not in globals():
-            import nvdiffrast.torch as dr
-            
         if not isinstance(envmap, dict):
             envmap = {'' : envmap}
         num_envmaps = len(envmap)
@@ -322,12 +314,11 @@ class PbrMeshRenderer:
                 
                 # PBR attributes
                 if isinstance(mesh, MeshWithVoxel):
-                    if 'grid_sample_3d' not in globals():
-                        from flex_gemm.ops.grid_sample import grid_sample_3d
+                    from ..utils.grid_sample import grid_sample_3d as _grid_sample_3d
                     mask = rast[..., -1:] > 0
                     xyz = dr.interpolate(vertices_orig, rast, faces)[0]
                     xyz = ((xyz - mesh.origin) / mesh.voxel_size).reshape(1, -1, 3)
-                    img = grid_sample_3d(
+                    img = _grid_sample_3d(
                         mesh.attrs,
                         torch.cat([torch.zeros_like(mesh.coords[..., :1]), mesh.coords], dim=-1),
                         mesh.voxel_shape,

trellis2/renderers/voxel_renderer.py

@@ -1,7 +1,6 @@
 import torch
 from easydict import EasyDict as edict
 from ..representations import Voxel
-from easydict import EasyDict as edict
 
 
 class VoxelRenderer:
@@ -29,7 +28,7 @@ class VoxelRenderer:
             colors_overwrite: torch.Tensor = None
         ) -> edict:
         """
-        Render the gausssian.
+        Render the voxel.
 
         Args:
             voxel (Voxel): Voxel representation.

trellis2/representations/mesh/base.py

@@ -1,8 +1,21 @@
 from typing import *
 import torch
+import torch.nn.functional as F_grid
+import numpy as np
 from ..voxel import Voxel
-import cumesh
-from flex_gemm.ops.grid_sample import grid_sample_3d
+
+from trellis2.backends import (
+    MeshBackend as _MeshBackendClass,
+    HAS_MESH as _HAS_MESH,
+    HAS_MPS as _HAS_MPS,
+    HAS_CUDA as _HAS_CUDA,
+    HAS_TRIMESH,
+    trimesh as _trimesh,
+    HAS_FAST_SIMPLIFICATION,
+    fast_simplification,
+    HAS_FLEX_GEMM,
+    _flex_grid_sample_3d as grid_sample_3d,
+)
 
 
 class Mesh:
@@ -32,11 +45,20 @@ class Mesh:
     def cpu(self):
         return self.to('cpu')
 
-    def fill_holes(self, max_hole_perimeter=3e-2):
-        vertices = self.vertices.cuda()
-        faces = self.faces.cuda()
+    def _gpu_device(self):
+        if _HAS_MPS:
+            return 'mps'
+        if _HAS_CUDA:
+            return 'cuda'
+        return None
 
-        mesh = cumesh.CuMesh()
+    def fill_holes(self, max_hole_perimeter=3e-2):
+        if _HAS_MESH and self._gpu_device():
+            dev = self._gpu_device()
+            vertices = self.vertices.to(dev)
+            faces = self.faces.to(dev)
+
+            mesh = _MeshBackendClass()
             mesh.init(vertices, faces)
             mesh.get_edges()
             mesh.get_boundary_info()
@@ -55,30 +77,76 @@ class Mesh:
 
             self.vertices = new_vertices.to(self.device)
             self.faces = new_faces.to(self.device)
+        elif HAS_TRIMESH:
+            tm = _trimesh.Trimesh(
+                vertices=self.vertices.cpu().numpy(),
+                faces=self.faces.cpu().numpy(),
+                process=False,
+            )
+            _trimesh.repair.fill_holes(tm)
+            self.vertices = torch.from_numpy(tm.vertices.astype(np.float32)).to(self.device)
+            self.faces = torch.from_numpy(tm.faces.astype(np.int32)).to(self.device)
+        else:
+            raise RuntimeError("No mesh repair backend available. Install trimesh or cumesh.")
 
     def remove_faces(self, face_mask: torch.Tensor):
-        vertices = self.vertices.cuda()
-        faces = self.faces.cuda()
+        if _HAS_MESH and self._gpu_device():
+            dev = self._gpu_device()
+            vertices = self.vertices.to(dev)
+            faces = self.faces.to(dev)
 
-        mesh = cumesh.CuMesh()
+            mesh = _MeshBackendClass()
             mesh.init(vertices, faces)
             mesh.remove_faces(face_mask)
             new_vertices, new_faces = mesh.read()
 
             self.vertices = new_vertices.to(self.device)
             self.faces = new_faces.to(self.device)
+        else:
+            # CPU fallback: mask out faces and re-index vertices
+            keep = ~face_mask.cpu()
+            new_faces = self.faces.cpu()[keep]
+            used_verts = torch.unique(new_faces)
+            vert_map = torch.full((self.vertices.shape[0],), -1, dtype=torch.long)
+            vert_map[used_verts] = torch.arange(len(used_verts))
+            self.vertices = self.vertices.cpu()[used_verts].to(self.device)
+            self.faces = vert_map[new_faces].int().to(self.device)
 
     def simplify(self, target=1000000, verbose: bool=False, options: dict={}):
-        vertices = self.vertices.cuda()
-        faces = self.faces.cuda()
+        if _HAS_MESH and self._gpu_device():
+            dev = self._gpu_device()
+            vertices = self.vertices.to(dev)
+            faces = self.faces.to(dev)
 
-        mesh = cumesh.CuMesh()
+            mesh = _MeshBackendClass()
             mesh.init(vertices, faces)
             mesh.simplify(target, verbose=verbose, options=options)
             new_vertices, new_faces = mesh.read()
 
             self.vertices = new_vertices.to(self.device)
             self.faces = new_faces.to(self.device)
+        elif HAS_FAST_SIMPLIFICATION:
+            verts_np = self.vertices.cpu().numpy().astype(np.float64)
+            faces_np = self.faces.cpu().numpy()
+            ratio = min(1.0, target / max(faces_np.shape[0], 1))
+            new_verts, new_faces = fast_simplification.simplify(
+                verts_np, faces_np, target_reduction=(1.0 - ratio)
+            )
+            self.vertices = torch.from_numpy(new_verts.astype(np.float32)).to(self.device)
+            self.faces = torch.from_numpy(new_faces.astype(np.int32)).to(self.device)
+        elif HAS_TRIMESH:
+            tm = _trimesh.Trimesh(
+                vertices=self.vertices.cpu().numpy(),
+                faces=self.faces.cpu().numpy(),
+                process=False,
+            )
+            if hasattr(tm, 'simplify_quadric_decimation'):
+                tm = tm.simplify_quadric_decimation(target)
+            self.vertices = torch.from_numpy(tm.vertices.astype(np.float32)).to(self.device)
+            self.faces = torch.from_numpy(tm.faces.astype(np.int32)).to(self.device)
+        else:
+            if verbose:
+                print("Warning: No simplification backend available, skipping.")
 
 
 class TextureFilterMode:
@@ -220,6 +288,7 @@ class MeshWithVoxel(Mesh, Voxel):
         )
         
     def query_attrs(self, xyz):
+        if HAS_FLEX_GEMM:
             grid = ((xyz - self.origin) / self.voxel_size).reshape(1, -1, 3)
             vertex_attrs = grid_sample_3d(
                 self.attrs,
@@ -230,5 +299,27 @@ class MeshWithVoxel(Mesh, Voxel):
             )[0]
             return vertex_attrs
 
+        # CPU/MPS fallback: build dense volume then F.grid_sample
+        C = self.attrs.shape[-1]
+        spatial = self.voxel_shape[2:]  # (D, H, W)
+        D, H, W = spatial
+        dense_vol = torch.zeros(1, C, D, H, W, dtype=self.attrs.dtype, device=self.attrs.device)
+        cx, cy, cz = self.coords[:, 0].long(), self.coords[:, 1].long(), self.coords[:, 2].long()
+        dense_vol[0, :, cx, cy, cz] = self.attrs.T
+
+        # Normalize query points to [-1, 1] for F.grid_sample
+        grid_pts = ((xyz - self.origin) / self.voxel_size)
+        # Normalize to [-1, 1]: grid_sample expects (z, y, x) ordering in the last dim
+        grid_pts_norm = torch.stack([
+            grid_pts[:, 2] / (W - 1) * 2 - 1,
+            grid_pts[:, 1] / (H - 1) * 2 - 1,
+            grid_pts[:, 0] / (D - 1) * 2 - 1,
+        ], dim=-1)
+        grid_pts_norm = grid_pts_norm.reshape(1, 1, 1, -1, 3)
+        sampled = F_grid.grid_sample(dense_vol, grid_pts_norm, mode='bilinear', align_corners=True, padding_mode='border')
+        # sampled: [1, C, 1, 1, N] -> [N, C]
+        vertex_attrs = sampled.reshape(C, -1).T
+        return vertex_attrs
+
     def query_vertex_attrs(self):
         return self.query_attrs(self.vertices)

trellis2/utils/grid_sample.py

@@ -0,0 +1,49 @@
+"""Unified 3D grid sampling — flex_gemm on Metal/CUDA, F.grid_sample fallback."""
+import torch
+import torch.nn.functional as F
+
+try:
+    from flex_gemm.ops.grid_sample import grid_sample_3d as _flex_grid_sample
+    _HAS_FLEX_GEMM = True
+except ImportError:
+    _HAS_FLEX_GEMM = False
+
+
+def grid_sample_3d(feats, coords, shape, grid, mode='trilinear'):
+    """Drop-in replacement for flex_gemm.ops.grid_sample.grid_sample_3d.
+
+    Args:
+        feats: [N, C] sparse voxel features
+        coords: [N, 4] voxel coordinates (batch_idx, x, y, z)
+        shape: torch.Size([B, C, D, H, W]) — sparse tensor shape
+        grid: [B, M, 3] query points in voxel space
+        mode: 'trilinear' (maps to F.grid_sample 'bilinear')
+    Returns:
+        [B*C, M] sampled features (matching flex_gemm output shape)
+    """
+    if _HAS_FLEX_GEMM:
+        return _flex_grid_sample(feats, coords, shape, grid, mode=mode)
+
+    # Dense volume fallback
+    B, C = shape[0], shape[1]
+    D, H, W = shape[2], shape[3], shape[4]
+    device = feats.device
+
+    dense_vol = torch.zeros(B, C, D, H, W, dtype=feats.dtype, device=device)
+    batch_idx = coords[:, 0].long()
+    cx, cy, cz = coords[:, 1].long(), coords[:, 2].long(), coords[:, 3].long()
+    dense_vol[batch_idx, :, cx, cy, cz] = feats
+
+    # Normalize grid to [-1, 1] for F.grid_sample (expects z,y,x order)
+    grid_norm = torch.stack([
+        grid[..., 2] / (W - 1) * 2 - 1,
+        grid[..., 1] / (H - 1) * 2 - 1,
+        grid[..., 0] / (D - 1) * 2 - 1,
+    ], dim=-1)
+    grid_norm = grid_norm.reshape(B, 1, 1, -1, 3)
+
+    sampled = F.grid_sample(dense_vol, grid_norm, mode='bilinear',
+                            align_corners=True, padding_mode='border')
+    # sampled: [B, C, 1, 1, M] -> reshape to match flex_gemm output
+    M = grid.shape[1]
+    return sampled.reshape(B * C, M)