diff --git a/.gitignore b/.gitignore index b7faf40..8d535c0 100644 --- a/.gitignore +++ b/.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/ diff --git a/README.md b/README.md index 1c0ad06..d94a714 100644 --- a/README.md +++ b/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 diff --git a/api_models.py b/api_models.py new file mode 100644 index 0000000..c6207c4 --- /dev/null +++ b/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 diff --git a/api_server.py b/api_server.py new file mode 100644 index 0000000..0849939 --- /dev/null +++ b/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() diff --git a/app_mlx.py b/app_mlx.py new file mode 100644 index 0000000..bbedc4d --- /dev/null +++ b/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() diff --git a/mlx_backend/__init__.py b/mlx_backend/__init__.py new file mode 100644 index 0000000..c303d7a --- /dev/null +++ b/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 diff --git a/mlx_backend/adapters.py b/mlx_backend/adapters.py new file mode 100644 index 0000000..c5637f8 --- /dev/null +++ b/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 diff --git a/mlx_backend/attention.py b/mlx_backend/attention.py new file mode 100644 index 0000000..bdb7a12 --- /dev/null +++ b/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) diff --git a/mlx_backend/convert.py b/mlx_backend/convert.py new file mode 100644 index 0000000..662614d --- /dev/null +++ b/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) diff --git a/mlx_backend/dinov3.py b/mlx_backend/dinov3.py new file mode 100644 index 0000000..c2056fb --- /dev/null +++ b/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 diff --git a/mlx_backend/flow_models.py b/mlx_backend/flow_models.py new file mode 100644 index 0000000..c9da0de --- /dev/null +++ b/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) diff --git a/mlx_backend/norm.py b/mlx_backend/norm.py new file mode 100644 index 0000000..7d7c7a4 --- /dev/null +++ b/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) diff --git a/mlx_backend/pipeline.py b/mlx_backend/pipeline.py new file mode 100644 index 0000000..58f8f99 --- /dev/null +++ b/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) diff --git a/mlx_backend/rope.py b/mlx_backend/rope.py new file mode 100644 index 0000000..f6ff19b --- /dev/null +++ b/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) diff --git a/mlx_backend/sparse_conv.py b/mlx_backend/sparse_conv.py new file mode 100644 index 0000000..4cc3a2b --- /dev/null +++ b/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)) diff --git a/mlx_backend/sparse_ops.py b/mlx_backend/sparse_ops.py new file mode 100644 index 0000000..3e7fa1f --- /dev/null +++ b/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 diff --git a/mlx_backend/sparse_tensor.py b/mlx_backend/sparse_tensor.py new file mode 100644 index 0000000..45a3631 --- /dev/null +++ b/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}") diff --git a/mlx_backend/structure_decoder.py b/mlx_backend/structure_decoder.py new file mode 100644 index 0000000..5b7bac5 --- /dev/null +++ b/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 diff --git a/mlx_backend/transformer_block.py b/mlx_backend/transformer_block.py new file mode 100644 index 0000000..9496fb3 --- /dev/null +++ b/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 diff --git a/mlx_backend/vae_blocks.py b/mlx_backend/vae_blocks.py new file mode 100644 index 0000000..ef649e5 --- /dev/null +++ b/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 diff --git a/mlx_backend/vae_decoders.py b/mlx_backend/vae_decoders.py new file mode 100644 index 0000000..991845a --- /dev/null +++ b/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) diff --git a/o-voxel/o_voxel/__init__.py b/o-voxel/o_voxel/__init__.py index 263b753..263656b 100644 --- a/o-voxel/o_voxel/__init__.py +++ b/o-voxel/o_voxel/__init__.py @@ -1,7 +1,17 @@ from . import ( convert, - io, postprocess, - rasterize, - serialize -) \ No newline at end of file + 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 \ No newline at end of file diff --git a/o-voxel/o_voxel/convert/__init__.py b/o-voxel/o_voxel/convert/__init__.py index f7b6d8c..6be6527 100644 --- a/o-voxel/o_voxel/convert/__init__.py +++ b/o-voxel/o_voxel/convert/__init__.py @@ -1,2 +1,6 @@ from .flexible_dual_grid import * -from .volumetic_attr import * \ No newline at end of file +try: + from .volumetic_attr import * +except ImportError: + # volumetic_attr requires _C extension + pass \ No newline at end of file diff --git a/o-voxel/o_voxel/convert/flexible_dual_grid.py b/o-voxel/o_voxel/convert/flexible_dual_grid.py index 7cf1397..c3cced7 100644 --- a/o-voxel/o_voxel/convert/flexible_dual_grid.py +++ b/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", @@ -11,7 +18,7 @@ __all__ = [ def _init_hashmap(grid_size, capacity, device): VOL = (grid_size[0] * grid_size[1] * grid_size[2]).item() - + # If the number of elements in the tensor is less than 2^32, use uint32 as the hashmap type, otherwise use uint64. if VOL < 2**32: hashmap_keys = torch.full((capacity,), torch.iinfo(torch.uint32).max, dtype=torch.uint32, device=device) @@ -21,10 +28,48 @@ def _init_hashmap(grid_size, capacity, device): raise ValueError(f"The spatial size is too large to fit in a hashmap. Get volumn {VOL} > 2^64.") hashmap_vals = torch.empty((capacity,), dtype=torch.uint32, device=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 - 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()) + 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,7 +284,10 @@ def flexible_dual_grid_to_mesh( torch.zeros((M * 4, 1), dtype=torch.int, device=coords.device), connected_voxel.reshape(-1, 3) ], dim=1) - connected_voxel_indices = _C.hashmap_lookup_3d_cuda(*hashmap, connected_voxel_hash_key, *grid_size.tolist()).reshape(M, 4).int() + 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) L = quad_indices.shape[0] @@ -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 diff --git a/o-voxel/o_voxel/postprocess.py b/o-voxel/o_voxel/postprocess.py index 1ce8227..c9fa5b0 100644 --- a/o-voxel/o_voxel/postprocess.py +++ b/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() - - # Initialize CUDA mesh handler - mesh = cumesh.CuMesh() + vertices = vertices.to(device) + faces = faces.to(device) + + # 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, @@ -180,8 +261,16 @@ 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,7 +371,14 @@ 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) - alpha_mode = 'OPAQUE' + # 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 mask_inv = (~mask).astype(np.uint8) @@ -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, diff --git a/o-voxel/o_voxel/postprocess_cpu.py b/o-voxel/o_voxel/postprocess_cpu.py new file mode 100644 index 0000000..e3c6e9c --- /dev/null +++ b/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 diff --git a/o-voxel/pyproject.toml b/o-voxel/pyproject.toml index 6b13d43..b0905fa 100644 --- a/o-voxel/pyproject.toml +++ b/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", ] diff --git a/o-voxel/setup.py b/o-voxel/setup.py index 91cb5ce..73ae873 100644 --- a/o-voxel/setup.py +++ b/o-voxel/setup.py @@ -1,26 +1,89 @@ 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: - cc_flag = [] -else: - archs = os.getenv("GPU_ARCHS", "native").split(";") - cc_flag = [f"--offload-arch={arch}" for arch in archs] + if not IS_HIP: + cc_flag = [] + 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", @@ -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, ) diff --git a/requirements_macos.txt b/requirements_macos.txt new file mode 100644 index 0000000..fe93378 --- /dev/null +++ b/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 ) +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 diff --git a/scripts/download_weights.py b/scripts/download_weights.py new file mode 100644 index 0000000..6250cdd --- /dev/null +++ b/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() diff --git a/scripts/start_server.sh b/scripts/start_server.sh new file mode 100755 index 0000000..f16f17d --- /dev/null +++ b/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" diff --git a/test_rast_parity.py b/test_rast_parity.py new file mode 100644 index 0000000..6029a92 --- /dev/null +++ b/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']) diff --git a/trellis2/backends.py b/trellis2/backends.py new file mode 100644 index 0000000..6ccae2e --- /dev/null +++ b/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 diff --git a/trellis2/modules/image_feature_extractor.py b/trellis2/modules/image_feature_extractor.py index c3cb515..8a635ac 100644 --- a/trellis2/modules/image_feature_extractor.py +++ b/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,24 +21,30 @@ 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: """ Extract features from the image. - + Args: image: A batch of images as a tensor of shape (B, C, H, W) or a list of PIL images. - + Returns: A tensor of shape (B, N, D) where N is the number of patches and D is the feature dimension. """ @@ -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) @@ -90,15 +105,15 @@ class DinoV3FeatureExtractor: ) return F.layer_norm(hidden_states, hidden_states.shape[-1:]) - + @torch.no_grad() def __call__(self, image: Union[torch.Tensor, List[Image.Image]]) -> torch.Tensor: """ Extract features from the image. - + Args: image: A batch of images as a tensor of shape (B, C, H, W) or a list of PIL images. - + Returns: A tensor of shape (B, N, D) where N is the number of patches and D is the feature dimension. """ @@ -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 diff --git a/trellis2/modules/sparse/attention/full_attn.py b/trellis2/modules/sparse/attention/full_attn.py index 4eb74d2..a44afa5 100644 --- a/trellis2/modules/sparse/attention/full_attn.py +++ b/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}") diff --git a/trellis2/modules/sparse/attention/windowed_attn.py b/trellis2/modules/sparse/attention/windowed_attn.py index 0430788..5b1e85a 100644 --- a/trellis2/modules/sparse/attention/windowed_attn.py +++ b/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] diff --git a/trellis2/modules/sparse/config.py b/trellis2/modules/sparse/config.py index a5f4d53..f6d1070 100644 --- a/trellis2/modules/sparse/config.py +++ b/trellis2/modules/sparse/config.py @@ -1,36 +1,61 @@ from typing import * +import platform +import sys -CONV = 'flex_gemm' +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 - + global CONV 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}") - + __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 diff --git a/trellis2/modules/sparse/conv/conv_pytorch.py b/trellis2/modules/sparse/conv/conv_pytorch.py new file mode 100644 index 0000000..71bd318 --- /dev/null +++ b/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) diff --git a/trellis2/pipelines/rembg/BiRefNet.py b/trellis2/pipelines/rembg/BiRefNet.py index c71a992..61056fc 100644 --- a/trellis2/pipelines/rembg/BiRefNet.py +++ b/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]), ] ) - - def to(self, device: str): + 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') + 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() diff --git a/trellis2/pipelines/trellis2_image_to_3d.py b/trellis2/pipelines/trellis2_image_to_3d.py index a7b84b9..465986b 100644 --- a/trellis2/pipelines/trellis2_image_to_3d.py +++ b/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,7 +231,8 @@ class Trellis2ImageTo3DPipeline(Pipeline): ).samples if self.low_vram: flow_model.cpu() - + _free_memory() + # Decode sparse structure latent decoder = self.models['sparse_structure_decoder'] if self.low_vram: @@ -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,13 +281,14 @@ 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) slat = slat * std + mean - + return slat - + def sample_shape_slat_cascade( self, lr_cond: dict, @@ -312,10 +327,11 @@ 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 - + # Upsample if self.low_vram: self.models['shape_slat_decoder'].to(self.device) @@ -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,11 +373,12 @@ 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) slat = slat * std + mean - + return slat, hr_resolution def decode_shape_slat( @@ -386,8 +404,9 @@ 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( self, cond: dict, @@ -424,11 +443,12 @@ 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) slat = slat * std + mean - + return slat def decode_tex_slat( @@ -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 ) - torch.cuda.empty_cache() + 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) diff --git a/trellis2/pipelines/trellis2_texturing.py b/trellis2/pipelines/trellis2_texturing.py index c184b5e..4fa5d64 100755 --- a/trellis2/pipelines/trellis2_texturing.py +++ b/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]), diff --git a/trellis2/renderers/mesh_renderer.py b/trellis2/renderers/mesh_renderer.py index e20efc5..98ee9fc 100644 --- a/trellis2/renderers/mesh_renderer.py +++ b/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, diff --git a/trellis2/renderers/pbr_mesh_renderer.py b/trellis2/renderers/pbr_mesh_renderer.py index b3b4655..def830c 100644 --- a/trellis2/renderers/pbr_mesh_renderer.py +++ b/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, diff --git a/trellis2/renderers/voxel_renderer.py b/trellis2/renderers/voxel_renderer.py index dfe28ad..8a48018 100644 --- a/trellis2/renderers/voxel_renderer.py +++ b/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. diff --git a/trellis2/representations/mesh/base.py b/trellis2/representations/mesh/base.py index b70e4cc..857efbf 100644 --- a/trellis2/representations/mesh/base.py +++ b/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: @@ -14,71 +27,126 @@ class Mesh: self.vertices = vertices.float() self.faces = faces.int() self.vertex_attrs = vertex_attrs - + @property def device(self): return self.vertices.device - + def to(self, device, non_blocking=False): return Mesh( self.vertices.to(device, non_blocking=non_blocking), self.faces.to(device, non_blocking=non_blocking), self.vertex_attrs.to(device, non_blocking=non_blocking) if self.vertex_attrs is not None else None, ) - + def cuda(self, non_blocking=False): return self.to('cuda', non_blocking=non_blocking) - + def cpu(self): return self.to('cpu') - + + def _gpu_device(self): + if _HAS_MPS: + return 'mps' + if _HAS_CUDA: + return 'cuda' + return None + def fill_holes(self, max_hole_perimeter=3e-2): - vertices = self.vertices.cuda() - faces = self.faces.cuda() - - mesh = cumesh.CuMesh() - mesh.init(vertices, faces) - mesh.get_edges() - mesh.get_boundary_info() - if mesh.num_boundaries == 0: - return - mesh.get_vertex_edge_adjacency() - mesh.get_vertex_boundary_adjacency() - mesh.get_manifold_boundary_adjacency() - mesh.read_manifold_boundary_adjacency() - mesh.get_boundary_connected_components() - mesh.get_boundary_loops() - if mesh.num_boundary_loops == 0: - return - mesh.fill_holes(max_hole_perimeter=max_hole_perimeter) - new_vertices, new_faces = mesh.read() - - self.vertices = new_vertices.to(self.device) - self.faces = new_faces.to(self.device) - + 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() + if mesh.num_boundaries == 0: + return + mesh.get_vertex_edge_adjacency() + mesh.get_vertex_boundary_adjacency() + mesh.get_manifold_boundary_adjacency() + mesh.read_manifold_boundary_adjacency() + mesh.get_boundary_connected_components() + mesh.get_boundary_loops() + if mesh.num_boundary_loops == 0: + return + mesh.fill_holes(max_hole_perimeter=max_hole_perimeter) + new_vertices, new_faces = mesh.read() + + 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() - - mesh = cumesh.CuMesh() - 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) - + 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.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() - - mesh = cumesh.CuMesh() - 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) + 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.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,15 +288,38 @@ class MeshWithVoxel(Mesh, Voxel): ) def query_attrs(self, xyz): - grid = ((xyz - self.origin) / self.voxel_size).reshape(1, -1, 3) - vertex_attrs = grid_sample_3d( - self.attrs, - torch.cat([torch.zeros_like(self.coords[..., :1]), self.coords], dim=-1), - self.voxel_shape, - grid, - mode='trilinear' - )[0] + if HAS_FLEX_GEMM: + grid = ((xyz - self.origin) / self.voxel_size).reshape(1, -1, 3) + vertex_attrs = grid_sample_3d( + self.attrs, + torch.cat([torch.zeros_like(self.coords[..., :1]), self.coords], dim=-1), + self.voxel_shape, + grid, + mode='trilinear' + )[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) diff --git a/trellis2/utils/grid_sample.py b/trellis2/utils/grid_sample.py new file mode 100644 index 0000000..894e94e --- /dev/null +++ b/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)