From f59e6e59819f50da6eb8d88f8a1c6c10c8693922 Mon Sep 17 00:00:00 2001 From: Jourloy Date: Fri, 17 Jul 2026 09:21:34 +0300 Subject: [PATCH] feat: add native Apple Silicon asset pipeline --- README.md | 94 +++- api_models.py | 35 -- api_server.py | 168 ------ app_mlx.py | 174 ------ mlx_backend/dinov3.py | 20 +- mlx_backend/pipeline.py | 64 ++- o-voxel/o_voxel/postprocess.py | 41 +- o-voxel/o_voxel/postprocess_cpu.py | 11 +- requirements_macos.txt | 50 +- requirements_macos_core.txt | 33 ++ scripts/download_weights.py | 118 ++-- scripts/generate_asset.py | 574 ++++++++++++++++++++ scripts/probe_macos.py | 118 ++++ scripts/setup_macos.sh | 118 ++++ scripts/validate_outputs.py | 28 + test_flex_gemm_integration.py | 31 +- tests/conftest.py | 6 + tests/test_backend_fallbacks.py | 114 ++++ tests/test_background_preprocess.py | 72 +++ tests/test_generate_asset.py | 128 +++++ tests/test_mlx_parity.py | 103 ++++ tests/test_model_revisions.py | 36 ++ trellis2/backends.py | 159 +++++- trellis2/gltf_validation.py | 151 +++++ trellis2/model_revisions.py | 34 ++ trellis2/models/__init__.py | 22 +- trellis2/modules/image_feature_extractor.py | 17 +- trellis2/modules/sparse/config.py | 73 ++- trellis2/pipelines/__init__.py | 9 +- trellis2/pipelines/base.py | 39 +- trellis2/pipelines/rembg/BiRefNet.py | 16 +- trellis2/pipelines/trellis2_image_to_3d.py | 33 +- trellis2/pipelines/trellis2_texturing.py | 23 +- trellis2/representations/mesh/base.py | 6 +- 34 files changed, 2161 insertions(+), 557 deletions(-) delete mode 100644 api_models.py delete mode 100644 api_server.py delete mode 100644 app_mlx.py create mode 100644 requirements_macos_core.txt mode change 100644 => 100755 scripts/download_weights.py create mode 100755 scripts/generate_asset.py create mode 100755 scripts/probe_macos.py create mode 100755 scripts/setup_macos.sh create mode 100755 scripts/validate_outputs.py create mode 100644 tests/conftest.py create mode 100644 tests/test_backend_fallbacks.py create mode 100644 tests/test_background_preprocess.py create mode 100644 tests/test_generate_asset.py create mode 100644 tests/test_mlx_parity.py create mode 100644 tests/test_model_revisions.py create mode 100644 trellis2/gltf_validation.py create mode 100644 trellis2/model_revisions.py diff --git a/README.md b/README.md index d94a714..6657445 100644 --- a/README.md +++ b/README.md @@ -17,7 +17,99 @@ https://github.com/user-attachments/assets/63b43a7e-acc7-4c81-a900-6da450527d8f ## 🍎 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. +This fork makes Apple Silicon a source-native backend instead of cloning and +patching a second TRELLIS checkout. The supported end-to-end path is PyTorch +MPS with capability-probed Metal extensions. MLX is included only as an +experimental parity backend; the upstream CUDA/Linux path remains available. + +The implementation incorporates the source-native backend and parity work from +[`trellis2-apple@6055b86`](https://github.com/pedronaugusto/trellis2-apple/commit/6055b868734af6e12769d229d90580e775fae9f0) +and the MPS CLI/fallback lessons from +[`trellis-mac@d58628f`](https://github.com/shivampkumar/trellis-mac/commit/d58628f4f5b9c3de8274cb110074154f4b31cef2). + +### macOS setup + +Requirements: Apple Silicon, Xcode, and Python 3.11. The setup script creates +`.venv`, installs the matching Xcode Metal Toolchain when necessary, tries +`torch==2.13.0` / `torchvision==0.28.0`, builds the pinned Metal extensions, +then runs `pip check`, MPS, SDPA, MLX, KDTree, and Metal probes. If the primary +pair fails its ABI probe, it rebuilds once with the prescribed +`torch==2.11.0` / `torchvision==0.26.0` fallback. + +```sh +git clone https://github.com/Jourloy/TRELLIS.2.git +cd TRELLIS.2 +bash scripts/setup_macos.sh +source .venv/bin/activate +``` + +The Metal sources are fixed to these commits: + +- `mtlgemm`: `867aec8234299a7fe1ede7f802c8debe5a939a82` +- `mtldiffrast`: `4668cd91cb6d27f5e264731f94a06841fbf7aab8` +- `mtlbvh`: `23f441c470ce1f537e1fd836f3ffb5b8245f7975` +- `mtlmesh`: `212079e55772cff3d648a21372392c37e0643f3b` + +`SKIP_METAL=1 bash scripts/setup_macos.sh` installs the slower controlled +fallback environment. No API or Gradio server is needed for the supported CLI. + +### Hugging Face access and offline cache + +TRELLIS.2-4B is public. DINOv3 and RMBG-2.0 are gated, so accept their terms +and authenticate once before downloading all pinned inputs: + +```sh +hf auth login +python scripts/download_weights.py \ + --cache-dir ~/.cache/trellis2/huggingface +python scripts/download_weights.py \ + --cache-dir ~/.cache/trellis2/huggingface --offline +``` + +Runtime revisions are fixed to TRELLIS.2-4B `af44b45`, DINOv3 `ea8dc28`, +RMBG-2.0 `5df4c9c`, and the external TRELLIS image decoder `25e0d31`. Pass the +same cache with `--cache-dir` and use `--offline` after the first download. + +### Reproducible generation + +```sh +python scripts/generate_asset.py input.png \ + --output-dir outputs/sample \ + --backend auto \ + --baker auto \ + --pipeline-type 512 \ + --seed 42 \ + --texture-size 1024 \ + --background auto \ + --pbr-decimation-target none \ + --cache-dir ~/.cache/trellis2/huggingface +``` + +`background=auto` runs official RMBG-2.0 for opaque images, while a prepared +RGBA image with non-opaque alpha bypasses background removal. `background=keep` +passes the image through unchanged. + +Every run writes: + +- `raw_full.glb`: untouched full-resolution geometry before PBR processing; +- `candidate_pbr.glb`: UV0 and native base-color, metallic, roughness, and + alpha data; +- `meta.json`: exact revisions, backend capabilities, seed, timings, hashes, + bounds, sizes, triangle counts, and every fallback attempt. + +There is no default triangle limit. PBR export first tries the complete mesh +with Metal, then full-resolution KDTree baking. If both fail on a very large +mesh, a separately recorded technical candidate near 200k faces is tried; +`raw_full.glb` is never simplified. Set `TRELLIS_DISABLE_METAL=1` to exercise +the pure-PyTorch sparse-convolution, SDPA, CPU mesh extraction, and KDTree +fallback path explicitly. + +### Dependency licenses + +The repository code is MIT, but model dependencies have their own terms. In +particular, review the DINOv3 license before use and note that RMBG-2.0 is +distributed under CC BY-NC 4.0. These terms are documented here and are not +enforced by an automatic runtime block. ## ✨ Features diff --git a/api_models.py b/api_models.py deleted file mode 100644 index c6207c4..0000000 --- a/api_models.py +++ /dev/null @@ -1,35 +0,0 @@ -"""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 deleted file mode 100644 index 2bb923a..0000000 --- a/api_server.py +++ /dev/null @@ -1,168 +0,0 @@ -""" -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 sys - -# Environment defaults — set before any torch/trellis imports -os.environ.setdefault("SPARSE_CONV_BACKEND", "flex_gemm") -os.environ.setdefault("ATTN_BACKEND", "sdpa") -os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1") # deform_conv2d in RMBG - -# Add o-voxel to path if present -_ovoxel = os.path.join(os.path.dirname(os.path.abspath(__file__)), "o-voxel") -if os.path.isdir(_ovoxel) and _ovoxel not in sys.path: - sys.path.insert(0, _ovoxel) -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 deleted file mode 100644 index bbedc4d..0000000 --- a/app_mlx.py +++ /dev/null @@ -1,174 +0,0 @@ -""" -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/dinov3.py b/mlx_backend/dinov3.py index c2056fb..b7a72b1 100644 --- a/mlx_backend/dinov3.py +++ b/mlx_backend/dinov3.py @@ -7,6 +7,7 @@ import mlx.core as mx import mlx.nn as nn import numpy as np from PIL import Image +from typing import Optional class MlxDINOv3PatchEmbed(nn.Module): @@ -226,14 +227,25 @@ class MlxDINOv3FeatureExtractor(nn.Module): return mx.array(np.stack(processed)) -def load_dinov3_from_hf(model_name: str = "facebook/dinov3-vitl16-pretrain-lvd1689m", - image_size: int = 512) -> MlxDINOv3FeatureExtractor: +def load_dinov3_from_hf( + model_name: str = "facebook/dinov3-vitl16-pretrain-lvd1689m", + image_size: int = 512, + *, + revision: Optional[str] = None, + cache_dir: Optional[str] = None, + local_files_only: bool = False, +) -> 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") + hub_kwargs = { + "revision": revision, + "cache_dir": cache_dir, + "local_files_only": local_files_only, + } + config_path = hf_hub_download(model_name, "config.json", **hub_kwargs) + weight_path = hf_hub_download(model_name, "model.safetensors", **hub_kwargs) with open(config_path) as f: config = json.load(f) diff --git a/mlx_backend/pipeline.py b/mlx_backend/pipeline.py index 58f8f99..2ba44c4 100644 --- a/mlx_backend/pipeline.py +++ b/mlx_backend/pipeline.py @@ -12,6 +12,12 @@ import json import time import logging +from trellis2.model_revisions import ( + DINOV3_REVISION, + RMBG_REVISION, + revision_for_repo, +) + import mlx.core as mx logger = logging.getLogger(__name__) @@ -30,23 +36,43 @@ from .adapters import ( ) -def _resolve_hf_path(rel_path: str) -> str: +def _resolve_hf_path( + rel_path: str, + *, + cache_dir: str = None, + local_files_only: bool = False, +) -> 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") + hub_kwargs = { + "revision": revision_for_repo(repo_id), + "cache_dir": cache_dir, + "local_files_only": local_files_only, + } + json_path = hf_hub_download(repo_id, f"{file_base}.json", **hub_kwargs) + hf_hub_download(repo_id, f"{file_base}.safetensors", **hub_kwargs) return json_path.rsplit('.json', 1)[0] -def _resolve_model_path(weights_path: str, rel_path: str) -> str: +def _resolve_model_path( + weights_path: str, + rel_path: str, + *, + cache_dir: str = None, + local_files_only: bool = False, +) -> 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) + return _resolve_hf_path( + rel_path, + cache_dir=cache_dir, + local_files_only=local_files_only, + ) def _load_mlx_flow_model(path: str, config: dict): @@ -164,7 +190,12 @@ def _get_loader(name: str, config: dict): raise ValueError(f"No loader for model '{name}' (type: {config['name']})") -def create_mlx_pipeline(weights_path: str = "weights/TRELLIS.2-4B"): +def create_mlx_pipeline( + weights_path: str = "weights/TRELLIS.2-4B", + *, + cache_dir: str = None, + local_files_only: bool = False, +): """Create upstream Trellis2ImageTo3DPipeline with MLX-backed models. All model compute runs in MLX. The upstream PT pipeline handles @@ -183,7 +214,12 @@ def create_mlx_pipeline(weights_path: str = "weights/TRELLIS.2-4B"): # Load all models with MLX adapters models = {} for name, rel_path in args['models'].items(): - path = _resolve_model_path(weights_path, rel_path) + path = _resolve_model_path( + weights_path, + rel_path, + cache_dir=cache_dir, + local_files_only=local_files_only, + ) with open(f"{path}.json") as f: model_config = json.load(f) @@ -219,11 +255,21 @@ def create_mlx_pipeline(weights_path: str = "weights/TRELLIS.2-4B"): # Image conditioning (MLX DINOv3) pipeline.image_cond_model = MlxImageCondAdapter( - load_dinov3_from_hf(args['image_cond_model']['args']['model_name']) + load_dinov3_from_hf( + args['image_cond_model']['args']['model_name'], + revision=DINOV3_REVISION, + cache_dir=cache_dir, + local_files_only=local_files_only, + ) ) # Background removal (PT — lightweight, used once) - pipeline.rembg_model = BiRefNet(**args['rembg_model']['args']) + pipeline.rembg_model = BiRefNet( + **args['rembg_model']['args'], + revision=RMBG_REVISION, + cache_dir=cache_dir, + local_files_only=local_files_only, + ) pipeline.low_vram = True pipeline._device = torch.device('cpu') diff --git a/o-voxel/o_voxel/postprocess.py b/o-voxel/o_voxel/postprocess.py index 1f16265..c70c879 100644 --- a/o-voxel/o_voxel/postprocess.py +++ b/o-voxel/o_voxel/postprocess.py @@ -8,27 +8,35 @@ import trimesh import trimesh.visual import platform +import os _HAS_DR = False _HAS_MESH = False _BACKEND = None dr = None +_METAL_DISABLED = os.environ.get('TRELLIS_DISABLE_METAL', '0') == '1' +_BACKEND_ERRORS = {} # Differentiable rasterizer — mtldiffrast (Metal) or nvdiffrast (CUDA) try: + if _METAL_DISABLED: + raise ImportError('Metal disabled by TRELLIS_DISABLE_METAL=1') import mtldiffrast.torch as dr _HAS_DR = True _BACKEND = 'metal' -except ImportError: +except (ImportError, RuntimeError, OSError) as exc: + _BACKEND_ERRORS['mtldiffrast'] = str(exc) try: import nvdiffrast.torch as dr _HAS_DR = True _BACKEND = 'cuda' - except ImportError: - pass + except (ImportError, RuntimeError, OSError) as cuda_exc: + _BACKEND_ERRORS['nvdiffrast'] = str(cuda_exc) # Mesh processing — cumesh auto-selects Metal/CUDA try: + if _METAL_DISABLED and platform.system() == 'Darwin': + raise ImportError('Metal disabled by TRELLIS_DISABLE_METAL=1') import cumesh _MeshBackend = cumesh.CuMesh _BVH = cumesh.cuBVH @@ -36,15 +44,18 @@ try: _HAS_MESH = True if _BACKEND is None: _BACKEND = 'metal' if platform.system() == 'Darwin' else 'cuda' -except ImportError: - pass +except (ImportError, RuntimeError, OSError) as exc: + _BACKEND_ERRORS['cumesh'] = str(exc) _HAS_GPU_DEPS = _HAS_DR and _HAS_MESH try: + if _METAL_DISABLED and platform.system() == 'Darwin': + raise ImportError('Metal disabled by TRELLIS_DISABLE_METAL=1') from flex_gemm.ops.grid_sample import grid_sample_3d as _flex_grid_sample_3d _HAS_FLEX_GEMM = True -except ImportError: +except (ImportError, RuntimeError, OSError) as exc: + _BACKEND_ERRORS['flex_gemm'] = str(exc) _HAS_FLEX_GEMM = False @@ -80,7 +91,7 @@ def to_glb( 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, + decimation_target: Optional[int] = 1000000, texture_size: int = 2048, remesh: bool = False, remesh_band: float = 1, @@ -105,7 +116,7 @@ def to_glb( aabb: (2, 3) tensor of minimum and maximum coordinates of the volume voxel_size: (3,) tensor of size of each voxel grid_size: (3,) tensor of number of voxels in each dimension - decimation_target: target number of vertices for mesh simplification + decimation_target: target face count, or None to preserve full resolution texture_size: size of the texture for baking remesh: whether to perform remeshing remesh_band: size of the remeshing band @@ -178,6 +189,12 @@ def to_glb( if verbose: print(f"Original mesh: {vertices.shape[0]} vertices, {faces.shape[0]} faces") + effective_decimation_target = ( + int(faces.shape[0]) if decimation_target is None else int(decimation_target) + ) + if effective_decimation_target <= 0: + effective_decimation_target = int(faces.shape[0]) + # Move data to GPU vertices = vertices.to(device) faces = faces.to(device) @@ -214,7 +231,7 @@ def to_glb( # --- Branch 1: Standard Pipeline (Simplification & Cleaning) --- if not remesh: # Step 1: Aggressive simplification (3x target) - mesh.simplify(decimation_target * 3, verbose=verbose) + mesh.simplify(effective_decimation_target * 3, verbose=verbose) if verbose: print(f"After inital simplification: {mesh.num_vertices} vertices, {mesh.num_faces} faces") @@ -227,7 +244,7 @@ def to_glb( print(f"After initial cleanup: {mesh.num_vertices} vertices, {mesh.num_faces} faces") # Step 3: Final simplification to target count - mesh.simplify(decimation_target, verbose=verbose) + mesh.simplify(effective_decimation_target, verbose=verbose) if verbose: print(f"After final simplification: {mesh.num_vertices} vertices, {mesh.num_faces} faces") @@ -271,7 +288,7 @@ def to_glb( print(f"After cleanup: {mesh.num_vertices} vertices, {mesh.num_faces} faces") # Simplify and clean the remeshed result - mesh.simplify(decimation_target, verbose=verbose) + mesh.simplify(effective_decimation_target, verbose=verbose) if verbose: print(f"After simplifying: {mesh.num_vertices} vertices, {mesh.num_faces} faces") @@ -424,4 +441,4 @@ def to_glb( if verbose: print("Done") - return textured_mesh \ No newline at end of file + return textured_mesh diff --git a/o-voxel/o_voxel/postprocess_cpu.py b/o-voxel/o_voxel/postprocess_cpu.py index de4f22d..ed6d795 100644 --- a/o-voxel/o_voxel/postprocess_cpu.py +++ b/o-voxel/o_voxel/postprocess_cpu.py @@ -218,7 +218,7 @@ def to_glb( 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, + decimation_target: Optional[int] = 1000000, texture_size: int = 2048, remesh: bool = False, remesh_band: float = 1, @@ -282,16 +282,19 @@ def to_glb( print(f"After hole filling: {len(tm.vertices)} vertices, {len(tm.faces)} faces") # --- Step 2: Simplification --- - if decimation_target < len(tm.faces): + effective_decimation_target = len(tm.faces) if decimation_target is None else int(decimation_target) + if effective_decimation_target <= 0: + effective_decimation_target = len(tm.faces) + if effective_decimation_target < len(tm.faces): if HAS_FAST_SIMPLIFICATION: - ratio = min(1.0, decimation_target / max(len(tm.faces), 1)) + ratio = min(1.0, effective_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) + tm = tm.simplify_quadric_decimation(effective_decimation_target) if verbose: print(f"After simplification: {len(tm.vertices)} vertices, {len(tm.faces)} faces") diff --git a/requirements_macos.txt b/requirements_macos.txt index ebda466..ba326de 100644 --- a/requirements_macos.txt +++ b/requirements_macos.txt @@ -1,45 +1,5 @@ -# Trellis2 MLX requirements (Apple Silicon) -# NO: flash_attn, spconv, torchsparse (CUDA-only) - -# Core -# torch >= 2.11.0 is required because mtlgemm's metal_context.mm calls -# at::mps::dispatch_sync_with_rethrow, which was promoted from -# at::native::mps:: to at::mps:: in PyTorch 2.11 (PR #167445, merged -# 2025-11-11). Earlier stable versions (2.6 – 2.10) only expose it under -# at::native::mps:: and won't link. -torch>=2.11.0 -torchvision>=0.26.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 +# Verified Apple Silicon runtime pair. scripts/setup_macos.sh retries the +# documented 2.11/0.26 ABI fallback only when the primary Metal build fails. +torch==2.13.0 +torchvision==0.28.0 +-r requirements_macos_core.txt diff --git a/requirements_macos_core.txt b/requirements_macos_core.txt new file mode 100644 index 0000000..9da95ea --- /dev/null +++ b/requirements_macos_core.txt @@ -0,0 +1,33 @@ +# Python runtime +transformers==4.57.6 +huggingface-hub==0.36.0 +safetensors==0.8.0 +pillow==12.3.0 +numpy==2.2.6 +scipy==1.17.1 + +# Inference and image processing +opencv-python-headless==4.12.0.88 +easydict==1.13 +einops==0.8.2 +kornia==0.8.3 +timm==1.0.28 +tqdm==4.68.4 +imageio==2.37.3 +imageio-ffmpeg==0.6.0 + +# Mesh and GLB processing +trimesh==4.12.2 +xatlas==0.0.11 +fast-simplification==0.1.13 +pygltflib==1.16.5 +plyfile==1.1.3 +zstandard==0.25.0 + +# Apple backend and diagnostics +mlx==0.32.0 +psutil==7.2.2 +pytest==9.1.1 + +# Upstream renderer utility, pinned instead of following main. +utils3d @ git+https://github.com/EasternJournalist/utils3d.git@9a4eb15e4021b67b12c460c7057d642626897ec8 diff --git a/scripts/download_weights.py b/scripts/download_weights.py old mode 100644 new mode 100755 index 6250cdd..4286811 --- a/scripts/download_weights.py +++ b/scripts/download_weights.py @@ -1,59 +1,93 @@ -"""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 -""" +#!/usr/bin/env python3 +"""Populate the dedicated Hugging Face cache with every pinned runtime input.""" from __future__ import annotations import argparse +import json import os +import sys +import time +from pathlib import Path + +# The Xet transport has produced non-resumable response-body decode failures +# on large model shards on macOS. Plain HTTP keeps .incomplete files in the +# selected HF cache and resumes them on the next attempt. +os.environ.setdefault("HF_HUB_DISABLE_XET", "1") from huggingface_hub import snapshot_download +from huggingface_hub.errors import GatedRepoError, HfHubHTTPError + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +from trellis2.model_revisions import MODEL_REVISIONS, TRELLIS_REPO -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") +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( - "--full-snapshot", - action="store_true", - help="Download the entire repo instead of just pipeline configs and checkpoints", + "--cache-dir", + type=Path, + default=Path.home() / ".cache" / "trellis2" / "huggingface", ) + parser.add_argument("--offline", action="store_true", help="verify that the pinned cache is complete") + parser.add_argument("--max-workers", type=int, default=2) + parser.add_argument("--retries", type=int, default=3) args = parser.parse_args() + if args.max_workers < 1: + parser.error("--max-workers must be at least 1") + if args.retries < 1: + parser.error("--retries must be at least 1") + cache_dir = args.cache_dir.expanduser().resolve() + cache_dir.mkdir(parents=True, exist_ok=True) - output_dir = os.path.abspath(args.output_dir) - os.makedirs(output_dir, exist_ok=True) + snapshots = {} + for repo_id, revision in MODEL_REVISIONS.items(): + for attempt in range(1, args.retries + 1): + try: + snapshots[repo_id] = snapshot_download( + repo_id=repo_id, + revision=revision, + cache_dir=str(cache_dir), + local_files_only=args.offline, + max_workers=args.max_workers, + ) + break + except GatedRepoError as exc: + raise RuntimeError( + f"access to {repo_id}@{revision} is gated; accept its terms and run " + "`hf auth login`, then retry" + ) from exc + except (OSError, RuntimeError, HfHubHTTPError) as exc: + if args.offline or attempt == args.retries: + raise RuntimeError( + f"failed to cache {repo_id}@{revision} after {attempt} attempt(s): {exc}" + ) from exc + delay = min(5, attempt * 2) + print( + f"Retrying {repo_id}@{revision} after transient download error " + f"({attempt}/{args.retries}): {exc}", + file=sys.stderr, + ) + time.sleep(delay) - 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}") + print( + json.dumps( + { + "cache_dir": str(cache_dir), + "primary_repo": TRELLIS_REPO, + "revisions": MODEL_REVISIONS, + "snapshots": snapshots, + "offline": args.offline, + "max_workers": args.max_workers, + }, + indent=2, + sort_keys=True, + ) + ) + return 0 if __name__ == "__main__": - main() + raise SystemExit(main()) diff --git a/scripts/generate_asset.py b/scripts/generate_asset.py new file mode 100755 index 0000000..0e5ca36 --- /dev/null +++ b/scripts/generate_asset.py @@ -0,0 +1,574 @@ +#!/usr/bin/env python3 +"""Generate reproducible full-resolution and PBR TRELLIS.2 assets.""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.metadata +import json +import os +import platform +import resource +import subprocess +import sys +import time +import traceback +from pathlib import Path +from typing import Any, Optional + +# MPS fallback must be configured before torch or any Metal extension imports. +os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1") +os.environ.setdefault("FLEX_GEMM_QUIET", "1") + +ROOT = Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +from trellis2.model_revisions import ( # noqa: E402 + DINOV3_REPO, + DINOV3_REVISION, + RMBG_REPO, + RMBG_REVISION, + SOURCE_REVISIONS, + TRELLIS_IMAGE_LARGE_REPO, + TRELLIS_IMAGE_LARGE_REVISION, + TRELLIS_REPO, + TRELLIS_REVISION, +) + +SAFETY_FACE_TARGET = 200_000 +WATCHDOG_SIGNATURES = ( + "non-zero size", + "BVH needs at least 8 triangles", + "kIOGPUCommandBufferCallbackErrorImpactingInteractivity", + "empty mesh", +) +OOM_SIGNATURES = ( + "out of memory", + "mps backend out of memory", + "failed to allocate", + "allocation failed", +) + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _write_json(path: Path, payload: dict[str, Any]) -> None: + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") + temporary.replace(path) + + +def _git_revision() -> Optional[str]: + result = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=ROOT, + text=True, + capture_output=True, + check=False, + ) + return result.stdout.strip() if result.returncode == 0 else None + + +def _peak_rss_bytes() -> int: + value = int(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss) + return value if sys.platform == "darwin" else value * 1024 + + +def _package_versions() -> dict[str, Optional[str]]: + names = ( + "torch", + "torchvision", + "transformers", + "huggingface-hub", + "mlx", + "flex-gemm", + "mtldiffrast", + "cumesh", + "trimesh", + "xatlas", + "fast-simplification", + ) + versions = {} + for name in names: + try: + versions[name] = importlib.metadata.version(name) + except importlib.metadata.PackageNotFoundError: + versions[name] = None + return versions + + +def _parse_target(value: str) -> Optional[int]: + if value.lower() in {"none", "full", "off", "0"}: + return None + target = int(value) + if target < 8: + raise argparse.ArgumentTypeError("decimation target must be at least 8 faces or 'none'") + return target + + +def _watchdog_message(error: BaseException) -> str: + detail = str(error) + if isinstance(error, MemoryError) or any( + signature in detail.lower() for signature in OOM_SIGNATURES + ): + return ( + "MPS ran out of memory. Close GPU-heavy apps, use one TRELLIS process, " + "and retry with --pipeline-type 512. Do not disable the MPS high-watermark " + f"guard. Original error: {detail}" + ) + if any(signature.lower() in detail.lower() for signature in WATCHDOG_SIGNATURES): + return ( + "The decoder produced an empty or watchdog-corrupted mesh. Close GPU-heavy apps, " + "retry pipeline 512, or run with TRELLIS_DISABLE_METAL=1 for the slower fallback. " + f"Original error: {detail}" + ) + return detail + + +def _raw_trimesh(mesh): + import numpy as np + import trimesh + + vertices = mesh.vertices.detach().cpu().numpy().astype(np.float32, copy=True) + faces = mesh.faces.detach().cpu().numpy().astype(np.int64, copy=True) + + # Match o_voxel.postprocess's asset-space conversion exactly. + converted = vertices.copy() + converted[:, 1] = vertices[:, 2] + converted[:, 2] = -vertices[:, 1] + raw = trimesh.Trimesh(vertices=converted, faces=faces, process=False) + # trimesh omits NORMAL from glTF when a process=False mesh has not yet + # materialized its cached vertex normals. Accessing the property computes + # them without merging vertices or changing the full-resolution topology. + _ = raw.vertex_normals + return raw + + +def _export_raw(mesh, path: Path) -> dict[str, Any]: + raw = _raw_trimesh(mesh) + raw.export(path) + return { + "vertices": int(len(raw.vertices)), + "triangles": int(len(raw.faces)), + "bounds_min": [float(value) for value in raw.bounds[0]], + "bounds_max": [float(value) for value in raw.bounds[1]], + "sha256": _sha256(path), + "bytes": path.stat().st_size, + } + + +def _metal_baker_available() -> tuple[bool, str]: + try: + from trellis2.backends import probe_metal_backends + from o_voxel import postprocess + + probes = probe_metal_backends() + failed_probes = { + name: result for name, result in probes.items() if not result.get("ok") + } + if failed_probes: + return False, json.dumps(failed_probes, sort_keys=True) + available = bool( + getattr(postprocess, "_HAS_DR", False) + and getattr(postprocess, "_HAS_MESH", False) + and getattr(postprocess, "_BACKEND", None) == "metal" + ) + if available: + return True, "" + errors = getattr(postprocess, "_BACKEND_ERRORS", {}) + return False, json.dumps(errors, sort_keys=True) if errors else "Metal backend is incomplete" + except (ImportError, RuntimeError, OSError) as exc: + return False, str(exc) + + +def _export_pbr(mesh, path: Path, *, baker: str, target: Optional[int], texture_size: int): + vertices = mesh.vertices.cpu() + faces = mesh.faces.cpu() + pre_simplified = False + + # cumesh builds its BVH before its own decimation pass. Reduce only the + # technical fallback candidate up front so the same oversized mesh cannot + # trip the Metal watchdog again. raw_full.glb is exported separately and + # remains untouched. + if target is not None and int(faces.shape[0]) > target: + import fast_simplification + import torch + + reduction = 1.0 - (target / int(faces.shape[0])) + simplified_vertices, simplified_faces = fast_simplification.simplify( + vertices.numpy(), + faces.numpy(), + target_reduction=reduction, + ) + vertices = torch.from_numpy(simplified_vertices).to(dtype=mesh.vertices.dtype) + faces = torch.from_numpy(simplified_faces).to(dtype=mesh.faces.dtype) + pre_simplified = True + + if baker == "metal": + available, reason = _metal_baker_available() + if not available: + raise RuntimeError(f"Metal baker unavailable: {reason}") + from o_voxel import postprocess + + exporter = postprocess.to_glb + elif baker == "kdtree": + from o_voxel import postprocess_cpu + + exporter = postprocess_cpu.to_glb + else: + raise ValueError(f"Unknown baker: {baker}") + + result = exporter( + vertices=vertices, + faces=faces, + attr_volume=mesh.attrs.cpu(), + coords=mesh.coords.cpu(), + attr_layout=mesh.layout, + voxel_size=mesh.voxel_size, + aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]], + decimation_target=target, + texture_size=texture_size, + verbose=True, + ) + result.export(path) + return result, pre_simplified + + +def _attempt_schedule(preferred: str, requested_target: Optional[int], raw_faces: int): + preferred_order = ["metal", "kdtree"] if preferred in {"auto", "metal"} else ["kdtree"] + targets = [requested_target] + if requested_target is None and raw_faces > SAFETY_FACE_TARGET: + targets.append(SAFETY_FACE_TARGET) + seen = set() + for target in targets: + for baker in preferred_order: + key = (baker, target) + if key not in seen: + seen.add(key) + yield baker, target + + +def _run_pbr_attempts( + mesh, + candidate_path: Path, + *, + preferred_baker: str, + requested_target: Optional[int], + texture_size: int, + export_fn=None, + on_attempt=None, +): + """Run the fixed Metal/KDTree/full/safety fallback order.""" + export_fn = export_fn or _export_pbr + raw_faces = int(mesh.faces.shape[0]) + attempts = [] + for baker, target in _attempt_schedule(preferred_baker, requested_target, raw_faces): + attempt = { + "baker": baker, + "target_faces": target, + "technical_safety_target": target == SAFETY_FACE_TARGET and raw_faces > SAFETY_FACE_TARGET, + } + attempt_started = time.perf_counter() + try: + exported, pre_simplified = export_fn( + mesh, + candidate_path, + baker=baker, + target=target, + texture_size=texture_size, + ) + attempt["status"] = "ok" + attempt["pre_simplified_before_bvh"] = pre_simplified + except Exception as exc: # each failure is recorded before the prescribed fallback + exported = None + attempt["status"] = "failed" + attempt["error_type"] = type(exc).__name__ + attempt["error"] = str(exc) + if candidate_path.exists(): + candidate_path.unlink() + attempt["duration_seconds"] = round(time.perf_counter() - attempt_started, 3) + attempts.append(attempt) + if on_attempt is not None: + on_attempt(attempt) + if exported is not None: + return exported, attempt, attempts + + errors = "; ".join( + f"{attempt['baker']}@{attempt['target_faces']}: {attempt.get('error', 'failed')}" + for attempt in attempts + ) + raise RuntimeError(f"All PBR export attempts failed: {errors}") + + +def _candidate_stats(path: Path, exported) -> dict[str, Any]: + import numpy as np + + vertices = np.asarray(exported.vertices) + faces = np.asarray(exported.faces) + return { + "vertices": int(len(vertices)), + "triangles": int(len(faces)), + "bounds_min": [float(value) for value in vertices.min(axis=0)], + "bounds_max": [float(value) for value in vertices.max(axis=0)], + "sha256": _sha256(path), + "bytes": path.stat().st_size, + } + + +def _load_pipeline(args): + import torch + + if args.backend == "mlx-experimental": + from huggingface_hub import snapshot_download + from mlx_backend.pipeline import create_mlx_pipeline + + snapshot = snapshot_download( + TRELLIS_REPO, + revision=TRELLIS_REVISION, + cache_dir=args.cache_dir, + local_files_only=args.offline, + ) + return create_mlx_pipeline( + snapshot, + cache_dir=args.cache_dir, + local_files_only=args.offline, + ), "mlx-experimental" + + if args.backend == "mps": + if platform.system() != "Darwin" or platform.machine() != "arm64": + raise RuntimeError("The mps backend requires an Apple Silicon Mac") + if not torch.backends.mps.is_available(): + raise RuntimeError("PyTorch MPS is not available in this Python environment") + elif args.backend == "cuda": + if not torch.cuda.is_available(): + raise RuntimeError("auto selected CUDA, but PyTorch CUDA is unavailable") + else: + raise RuntimeError(f"Unsupported resolved backend: {args.backend}") + + from trellis2.pipelines import Trellis2ImageTo3DPipeline + + pipeline = Trellis2ImageTo3DPipeline.from_pretrained( + TRELLIS_REPO, + revision=TRELLIS_REVISION, + cache_dir=args.cache_dir, + local_files_only=args.offline, + ) + if args.backend == "mps": + pipeline.to(torch.device("mps")) + return pipeline, "mps" + if args.backend == "cuda": + pipeline.cuda() + return pipeline, "cuda" + raise AssertionError("resolved backend validation fell through") + + +def _resolve_backend(requested: str) -> str: + if requested != "auto": + return requested + if platform.system() == "Darwin" and platform.machine() == "arm64": + return "mps" + + # Keep the official Linux/CUDA route intact without adding a new public + # CLI choice: --backend auto resolves to CUDA on supported non-Mac hosts. + import torch + + if torch.cuda.is_available(): + return "cuda" + raise RuntimeError("auto found neither Apple MPS nor NVIDIA CUDA") + + +def _parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("image", type=Path) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--backend", choices=("auto", "mps", "mlx-experimental"), default="auto") + parser.add_argument("--baker", choices=("auto", "metal", "kdtree"), default="auto") + parser.add_argument("--pipeline-type", choices=("512", "1024", "1024_cascade"), default="512") + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--steps", type=int) + parser.add_argument("--texture-size", type=int, choices=(512, 1024, 2048), default=1024) + parser.add_argument("--background", choices=("auto", "keep"), default="auto") + parser.add_argument("--pbr-decimation-target", type=_parse_target, default=None, metavar="none|N") + parser.add_argument("--cache-dir", type=Path) + parser.add_argument("--offline", action="store_true") + parser.add_argument("--force", action="store_true") + return parser.parse_args() + + +def main() -> int: + args = _parse_args() + requested_backend = args.backend + args.image = args.image.expanduser().resolve() + args.output_dir = args.output_dir.expanduser().resolve() + args.cache_dir = str(args.cache_dir.expanduser().resolve()) if args.cache_dir else None + + if not args.image.is_file(): + raise SystemExit(f"Input image not found: {args.image}") + if args.offline: + os.environ["HF_HUB_OFFLINE"] = "1" + args.backend = _resolve_backend(args.backend) + + args.output_dir.mkdir(parents=True, exist_ok=True) + raw_path = args.output_dir / "raw_full.glb" + candidate_path = args.output_dir / "candidate_pbr.glb" + metadata_path = args.output_dir / "meta.json" + existing = [path for path in (raw_path, candidate_path, metadata_path) if path.exists()] + if existing and not args.force: + raise SystemExit("Refusing to overwrite existing outputs: " + ", ".join(map(str, existing))) + if args.force: + for path in existing: + path.unlink() + + from PIL import Image + + with Image.open(args.image) as opened: + image = opened.copy() + input_size = list(opened.size) + input_mode = opened.mode + + metadata: dict[str, Any] = { + "schema_version": 1, + "status": "running", + "input": { + "path": str(args.image), + "sha256": _sha256(args.image), + "size": input_size, + "mode": input_mode, + }, + "configuration": { + "backend": requested_backend, + "resolved_backend": args.backend, + "preferred_baker": args.baker, + "pipeline_type": args.pipeline_type, + "seed": args.seed, + "steps": args.steps, + "texture_size": args.texture_size, + "background": args.background, + "pbr_decimation_target": args.pbr_decimation_target, + "offline": args.offline, + "cache_dir": args.cache_dir, + }, + "revisions": { + "code": _git_revision(), + TRELLIS_REPO: TRELLIS_REVISION, + TRELLIS_IMAGE_LARGE_REPO: TRELLIS_IMAGE_LARGE_REVISION, + DINOV3_REPO: DINOV3_REVISION, + RMBG_REPO: RMBG_REVISION, + }, + "source_revisions": SOURCE_REVISIONS, + "platform": { + "system": platform.system(), + "release": platform.release(), + "machine": platform.machine(), + "python": platform.python_version(), + }, + "packages": _package_versions(), + "timings_seconds": {}, + "pbr_attempts": [], + } + _write_json(metadata_path, metadata) + + started = time.perf_counter() + try: + load_started = time.perf_counter() + pipeline, resolved_backend = _load_pipeline(args) + metadata["configuration"]["resolved_backend"] = resolved_backend + try: + from trellis2.backends import backend_report + + metadata["backend_capabilities"] = backend_report() + except (ImportError, RuntimeError, OSError) as exc: + metadata["backend_capabilities"] = {"probe_error": str(exc)} + metadata["timings_seconds"]["pipeline_load"] = round(time.perf_counter() - load_started, 3) + _write_json(metadata_path, metadata) + + sampler_overrides = {"steps": args.steps} if args.steps is not None else {} + generation_started = time.perf_counter() + try: + outputs = pipeline.run( + image, + seed=args.seed, + pipeline_type=args.pipeline_type, + preprocess_image=args.background == "auto", + sparse_structure_sampler_params=sampler_overrides, + shape_slat_sampler_params=sampler_overrides, + tex_slat_sampler_params=sampler_overrides, + ) + except (IndexError, AssertionError, RuntimeError, MemoryError) as exc: + raise RuntimeError(_watchdog_message(exc)) from exc + metadata["timings_seconds"]["generation"] = round(time.perf_counter() - generation_started, 3) + + mesh = outputs[0] if isinstance(outputs, list) else outputs + if mesh.vertices.shape[0] == 0 or mesh.faces.shape[0] < 8: + raise RuntimeError(_watchdog_message(RuntimeError("empty mesh"))) + + raw_started = time.perf_counter() + metadata["raw_full"] = _export_raw(mesh, raw_path) + metadata["timings_seconds"]["raw_export"] = round(time.perf_counter() - raw_started, 3) + _write_json(metadata_path, metadata) + + pbr_started = time.perf_counter() + raw_faces = int(mesh.faces.shape[0]) + + def record_attempt(attempt): + metadata["pbr_attempts"].append(attempt) + _write_json(metadata_path, metadata) + + exported, chosen, _ = _run_pbr_attempts( + mesh, + candidate_path, + preferred_baker=args.baker, + requested_target=args.pbr_decimation_target, + texture_size=args.texture_size, + on_attempt=record_attempt, + ) + + metadata["candidate_pbr"] = _candidate_stats(candidate_path, exported) + metadata["candidate_pbr"].update(chosen) + metadata["candidate_pbr"]["technical_decimation"] = bool( + chosen["target_faces"] is not None and chosen["target_faces"] < raw_faces + ) + metadata["timings_seconds"]["pbr_export"] = round(time.perf_counter() - pbr_started, 3) + from trellis2.gltf_validation import validate_output_pair + + validation_started = time.perf_counter() + metadata["validation"] = validate_output_pair(raw_path, candidate_path) + metadata["timings_seconds"]["validation"] = round( + time.perf_counter() - validation_started, 3 + ) + metadata["timings_seconds"]["total"] = round(time.perf_counter() - started, 3) + metadata["peak_rss_bytes"] = _peak_rss_bytes() + metadata["status"] = "ok" + _write_json(metadata_path, metadata) + + print(f"raw_full.glb: {metadata['raw_full']['triangles']:,} triangles") + print( + "candidate_pbr.glb: " + f"{metadata['candidate_pbr']['triangles']:,} triangles via {chosen['baker']}" + ) + print(f"metadata: {metadata_path}") + return 0 + except Exception as exc: + metadata["status"] = "failed" + metadata["error"] = { + "type": type(exc).__name__, + "message": str(exc), + "traceback": traceback.format_exc(), + } + metadata["timings_seconds"]["total"] = round(time.perf_counter() - started, 3) + metadata["peak_rss_bytes"] = _peak_rss_bytes() + _write_json(metadata_path, metadata) + print(f"ERROR: {exc}", file=sys.stderr) + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/probe_macos.py b/scripts/probe_macos.py new file mode 100755 index 0000000..8c6f9e8 --- /dev/null +++ b/scripts/probe_macos.py @@ -0,0 +1,118 @@ +#!/usr/bin/env python3 +"""Probe the supported MPS, Metal, SDPA, MLX, and fallback paths.""" + +from __future__ import annotations + +import argparse +import json +import os +import platform +import subprocess +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) +sys.path.insert(0, str(ROOT / "o-voxel")) + + +def _metal_compiler() -> dict: + result = subprocess.run( + ["xcrun", "--find", "metal"], + text=True, + capture_output=True, + check=False, + ) + return { + "available": result.returncode == 0, + "path": result.stdout.strip() or None, + "error": result.stderr.strip() or None, + } + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--require-metal", action="store_true") + args = parser.parse_args() + + import mlx.core as mx + import numpy as np + import torch + import torch.nn.functional as torch_f + from scipy.spatial import cKDTree + + report = { + "platform": {"system": platform.system(), "machine": platform.machine()}, + "torch": torch.__version__, + "mps_built": torch.backends.mps.is_built(), + "mps_available": torch.backends.mps.is_available(), + "metal_compiler": _metal_compiler(), + } + failures = [] + + if not report["mps_available"]: + failures.append("PyTorch MPS is unavailable") + else: + device = torch.device("mps") + left = torch.arange(16, dtype=torch.float32, device=device).reshape(4, 4) + report["mps_matmul_sum"] = float((left @ left.T).sum().cpu()) + query = torch.randn(2, 2, 8, 16, device=device) + sdpa = torch_f.scaled_dot_product_attention(query, query, query) + report["sdpa_shape"] = list(sdpa.shape) + + mlx_matrix = mx.arange(16, dtype=mx.float32).reshape(4, 4) + mlx_value = mx.sum(mlx_matrix @ mlx_matrix.T) + mx.eval(mlx_value) + report["mlx_matmul_sum"] = float(mlx_value.item()) + + tree = cKDTree(np.array([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]])) + distance, index = tree.query([[0.1, 0.1, 0.1]], k=1) + report["kdtree_probe"] = {"distance": float(distance[0]), "index": int(index[0])} + + try: + from trellis2.backends import backend_report + + report["backends"] = backend_report() + except (ImportError, RuntimeError, OSError) as exc: + report["backends"] = {"probe_error": str(exc)} + + try: + from trellis2.modules.sparse import config as sparse_config + + report["sparse_backends"] = { + "convolution": sparse_config.CONV, + "attention": sparse_config.ATTN, + "metal_attention_parity": ( + sparse_config.probe_flex_gemm_sparse_attention_on_mps() + ), + } + except (ImportError, RuntimeError, OSError) as exc: + report["sparse_backends"] = {"probe_error": str(exc)} + + if args.require_metal: + backends = report.get("backends", {}) + capability_probes = backends.get("capability_probes", {}) + sparse_backends = report.get("sparse_backends", {}) + if not report["metal_compiler"]["available"]: + failures.append("Xcode Metal compiler is unavailable") + if backends.get("rasterizer") != "metal": + failures.append("mtldiffrast Metal rasterizer failed capability probe") + if backends.get("mesh") != "metal": + failures.append("mtlmesh/mtlbvh failed capability probe") + if not backends.get("flex_gemm"): + failures.append("mtlgemm/flex_gemm failed capability probe") + if sparse_backends.get("convolution") != "flex_gemm": + failures.append("flex_gemm sparse-convolution capability probe failed") + if not capability_probes.get("rasterizer", {}).get("ok"): + failures.append("mtldiffrast raster capability probe failed") + if not capability_probes.get("mesh_bvh", {}).get("ok"): + failures.append("mtlmesh/mtlbvh capability probe failed") + + report["ok"] = not failures + report["failures"] = failures + print(json.dumps(report, indent=2, sort_keys=True)) + return 0 if not failures else 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/setup_macos.sh b/scripts/setup_macos.sh new file mode 100755 index 0000000..e31ac89 --- /dev/null +++ b/scripts/setup_macos.sh @@ -0,0 +1,118 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +PYTHON_BIN="${PYTHON_BIN:-python3.11}" +VENV_DIR="${VENV_DIR:-$ROOT/.venv}" +DEPS_DIR="${TRELLIS_DEPS_DIR:-$HOME/.cache/trellis2/source-deps}" +PRIMARY_TORCH="2.13.0" +PRIMARY_TORCHVISION="0.28.0" +FALLBACK_TORCH="2.11.0" +FALLBACK_TORCHVISION="0.26.0" + +MTLGEMM_SHA="867aec8234299a7fe1ede7f802c8debe5a939a82" +MTLDIFFRAST_SHA="4668cd91cb6d27f5e264731f94a06841fbf7aab8" +MTLBVH_SHA="23f441c470ce1f537e1fd836f3ffb5b8245f7975" +MTLMESH_SHA="212079e55772cff3d648a21372392c37e0643f3b" + +if [[ "$(uname -s)" != "Darwin" || "$(uname -m)" != "arm64" ]]; then + echo "ERROR: this installer supports Apple Silicon macOS only" >&2 + exit 2 +fi +if ! command -v "$PYTHON_BIN" >/dev/null 2>&1; then + echo "ERROR: Python 3.11 is required (set PYTHON_BIN if it is installed elsewhere)" >&2 + exit 2 +fi +if ! "$PYTHON_BIN" -c 'import sys; raise SystemExit(sys.version_info[:2] != (3, 11))'; then + echo "ERROR: $PYTHON_BIN is not Python 3.11" >&2 + exit 2 +fi + +if [[ "${SKIP_METAL:-0}" != "1" ]] && ! xcrun --find metal >/dev/null 2>&1; then + echo "Metal Toolchain is missing; downloading the matching Xcode component..." + xcodebuild -downloadComponent MetalToolchain +fi + +if [[ ! -d "$VENV_DIR" ]]; then + "$PYTHON_BIN" -m venv "$VENV_DIR" +fi + +PYTHON="$VENV_DIR/bin/python" +PIP=("$PYTHON" -m pip) +"${PIP[@]}" install --upgrade 'pip==26.0.1' 'setuptools==80.9.0' 'wheel==0.46.3' 'ninja==1.13.0' 'pybind11==3.0.1' +"${PIP[@]}" install -r "$ROOT/requirements_macos_core.txt" + +install_torch_pair() { + local torch_version="$1" + local torchvision_version="$2" + "${PIP[@]}" uninstall -y torch torchvision mtldiffrast mtlbvh cumesh flex_gemm o_voxel >/dev/null 2>&1 || true + "${PIP[@]}" install "torch==$torch_version" "torchvision==$torchvision_version" +} + +checkout_dependency() { + local name="$1" + local url="$2" + local revision="$3" + local destination="$DEPS_DIR/$name-$revision" + local attempt + if [[ -d "$destination/.git" ]] && [[ "$(git -C "$destination" rev-parse HEAD)" == "$revision" ]]; then + DEP_PATH="$destination" + return 0 + fi + mkdir -p "$DEPS_DIR" + for attempt in 1 2 3; do + rm -rf "$destination" + if git -c http.version=HTTP/1.1 clone --filter=blob:none "$url" "$destination" \ + && git -C "$destination" -c http.version=HTTP/1.1 fetch --depth 1 origin "$revision" \ + && git -C "$destination" switch --detach "$revision" \ + && git -C "$destination" -c http.version=HTTP/1.1 submodule update --init --recursive; then + DEP_PATH="$destination" + return 0 + fi + echo "Dependency checkout failed ($name, attempt $attempt/3); retrying..." >&2 + done + echo "ERROR: failed to checkout pinned dependency $name@$revision" >&2 + return 1 +} + +install_metal_stack() { + export MACOSX_DEPLOYMENT_TARGET="${MACOSX_DEPLOYMENT_TARGET:-12.0}" + checkout_dependency mtlbvh https://github.com/pedronaugusto/mtlbvh.git "$MTLBVH_SHA" + "${PIP[@]}" install --no-build-isolation "$DEP_PATH" + checkout_dependency mtldiffrast https://github.com/pedronaugusto/mtldiffrast.git "$MTLDIFFRAST_SHA" + "${PIP[@]}" install --no-build-isolation "$DEP_PATH" + checkout_dependency mtlmesh https://github.com/pedronaugusto/mtlmesh.git "$MTLMESH_SHA" + "${PIP[@]}" install --no-build-isolation "$DEP_PATH" + checkout_dependency mtlgemm https://github.com/pedronaugusto/mtlgemm.git "$MTLGEMM_SHA" + "${PIP[@]}" install --no-build-isolation "$DEP_PATH" + BUILD_TARGET=cpu "${PIP[@]}" install --no-build-isolation -e "$ROOT/o-voxel" +} + +install_torch_pair "$PRIMARY_TORCH" "$PRIMARY_TORCHVISION" +resolved_torch="$PRIMARY_TORCH" +resolved_torchvision="$PRIMARY_TORCHVISION" + +if [[ "${SKIP_METAL:-0}" == "1" ]]; then + BUILD_TARGET=cpu "${PIP[@]}" install --no-build-isolation -e "$ROOT/o-voxel" +else + install_metal_stack + if ! "$PYTHON" "$ROOT/scripts/probe_macos.py" --require-metal; then + echo "Primary PyTorch pair failed the Metal build/probe; retrying the prescribed ABI fallback..." >&2 + install_torch_pair "$FALLBACK_TORCH" "$FALLBACK_TORCHVISION" + install_metal_stack + "$PYTHON" "$ROOT/scripts/probe_macos.py" --require-metal + resolved_torch="$FALLBACK_TORCH" + resolved_torchvision="$FALLBACK_TORCHVISION" + fi +fi + +"${PIP[@]}" check +"$PYTHON" "$ROOT/scripts/probe_macos.py" + +echo +echo "TRELLIS.2 Apple environment is ready." +echo " source $VENV_DIR/bin/activate" +echo " torch==$resolved_torch torchvision==$resolved_torchvision" +if ! "$PYTHON" -c 'from huggingface_hub import get_token; raise SystemExit(not bool(get_token()))' 2>/dev/null; then + echo " Hugging Face auth is missing: run '$VENV_DIR/bin/hf auth login' before downloading gated weights." +fi diff --git a/scripts/validate_outputs.py b/scripts/validate_outputs.py new file mode 100755 index 0000000..704f440 --- /dev/null +++ b/scripts/validate_outputs.py @@ -0,0 +1,28 @@ +#!/usr/bin/env python3 +"""Validate a raw_full.glb/candidate_pbr.glb pair and print JSON.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +from trellis2.gltf_validation import validate_output_pair + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("raw_full", type=Path) + parser.add_argument("candidate_pbr", type=Path) + args = parser.parse_args() + report = validate_output_pair(args.raw_full.resolve(), args.candidate_pbr.resolve()) + print(json.dumps(report, indent=2, sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/test_flex_gemm_integration.py b/test_flex_gemm_integration.py index 281a1cc..f942c5f 100644 --- a/test_flex_gemm_integration.py +++ b/test_flex_gemm_integration.py @@ -5,10 +5,10 @@ runs a SparseConv3d through the flex_gemm backend, and feeds the result through a LayerNorm — the exact path that originally crashed with "Passed CPU tensor to MPS op" before mtlgemm's device-routing fix. -Exercises both Algorithm.IMPLICIT_GEMM (default dense kernel) and -Algorithm.MASKED_IMPLICIT_GEMM (the masked kernel that landed in mtlgemm -round 2). Both should produce numerically equivalent output and pass the -LayerNorm hand-off without crashing. +Exercises Algorithm.IMPLICIT_GEMM, Algorithm.MASKED_IMPLICIT_GEMM, and the +production default Algorithm.MASKED_IMPLICIT_GEMM_SPLITK. All should produce +numerically equivalent output and pass the LayerNorm hand-off without +crashing. Note: this script intentionally uses torch.nn.functional.layer_norm rather than trellis2's SparseLayerNorm wrapper, because that wrapper currently calls @@ -26,7 +26,8 @@ import torch assert torch.backends.mps.is_available(), "This test needs MPS" from trellis2.modules.sparse import SparseTensor, SparseConv3d -from flex_gemm.ops.spconv import Algorithm, set_algorithm +from trellis2.modules.sparse.conv import config as conv_config +from flex_gemm.ops.spconv import Algorithm device = "mps" dtype = torch.float16 @@ -57,8 +58,18 @@ if conv.bias is not None: print(f"Conv weight device: {conv.weight.device}, dtype: {conv.weight.dtype}") def _run_with_algo(algo, label): - set_algorithm(algo) - out = conv(x) + # conv_flex_gemm sets flex_gemm's global algorithm from this config on + # every forward call, so changing flex_gemm directly would be overwritten. + conv_config.FLEX_GEMM_ALGO = algo + # Neighbor-cache layouts are algorithm-specific. Use a fresh sparse tensor + # so this parity test cannot feed one algorithm another algorithm's cache. + run_x = SparseTensor( + feats=feats, + coords=coords, + shape=torch.Size([1, ch]), + spatial_shape=[res, res, res], + ) + out = conv(run_x) assert out.feats.device.type == "mps", f"FAIL [{label}]: SparseConv3d on {out.feats.device}, expected mps" assert out.feats.dtype == dtype, f"FAIL [{label}]: SparseConv3d dtype {out.feats.dtype}, expected {dtype}" # LayerNorm hand-off — the original crash site @@ -73,13 +84,17 @@ def _run_with_algo(algo, label): print() y_dense = _run_with_algo(Algorithm.IMPLICIT_GEMM, "IMPLICIT_GEMM (dense)") y_masked = _run_with_algo(Algorithm.MASKED_IMPLICIT_GEMM, "MASKED_IMPLICIT_GEMM") +y_splitk = _run_with_algo(Algorithm.MASKED_IMPLICIT_GEMM_SPLITK, "MASKED_IMPLICIT_GEMM_SPLITK") # Numerical parity between the two algorithms — same inputs, equivalent output. diff = (y_dense - y_masked).abs().max().item() +splitk_diff = (y_dense - y_splitk).abs().max().item() parity_tol = 2e-2 # fp16 — masked reduces in a different order assert diff <= parity_tol, f"FAIL: dense vs masked diff {diff:.4e} > tol {parity_tol:.4e}" +assert splitk_diff <= parity_tol, f"FAIL: dense vs split-k diff {splitk_diff:.4e} > tol {parity_tol:.4e}" print(f" parity dense vs masked: max_diff={diff:.4e} (tol={parity_tol})") +print(f" parity dense vs split-k: max_diff={splitk_diff:.4e} (tol={parity_tol})") print() print("PASS — trellis2-apple SparseConv3d + LayerNorm runs end-to-end on MPS via flex_gemm") -print(" Both IMPLICIT_GEMM and MASKED_IMPLICIT_GEMM produce equivalent output.") +print(" Dense, masked, and production split-k algorithms produce equivalent output.") diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..8a6c0f3 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,6 @@ +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) +sys.path.insert(0, str(ROOT / "o-voxel")) diff --git a/tests/test_backend_fallbacks.py b/tests/test_backend_fallbacks.py new file mode 100644 index 0000000..bc4cbfe --- /dev/null +++ b/tests/test_backend_fallbacks.py @@ -0,0 +1,114 @@ +import json +import os +import platform +import subprocess +import sys +from pathlib import Path + +import pytest + +ROOT = Path(__file__).resolve().parents[1] + + +def test_disable_metal_selects_controlled_pytorch_fallback(): + env = os.environ.copy() + env.update( + { + "TRELLIS_DISABLE_METAL": "1", + "PYTHONPATH": str(ROOT), + "SPARSE_CONV_BACKEND": "pytorch", + "SPARSE_ATTN_BACKEND": "sdpa", + } + ) + code = """ +import json +from trellis2.modules.sparse import config +from trellis2.backends import backend_report +print('REPORT=' + json.dumps({ + 'conv': config.CONV, + 'attention': config.ATTN, + 'backend': backend_report(), +}, sort_keys=True)) +""" + completed = subprocess.run( + [sys.executable, "-c", code], + cwd=ROOT, + env=env, + text=True, + capture_output=True, + check=True, + ) + line = next(line for line in completed.stdout.splitlines() if line.startswith("REPORT=")) + report = json.loads(line.removeprefix("REPORT=")) + assert report["conv"] == "pytorch" + assert report["attention"] == "sdpa" + assert report["backend"]["metal_disabled"] is True + assert report["backend"]["rasterizer"] is None + assert report["backend"]["mesh"] is None + + +@pytest.mark.skipif(platform.system() != "Darwin", reason="Metal probe requires macOS") +def test_metal_sparse_attention_has_its_own_numerical_probe(): + import torch + + if not torch.backends.mps.is_available(): + pytest.skip("MPS is unavailable") + + from trellis2.modules.sparse import config + + assert config.probe_flex_gemm_sparse_attention_on_mps() + assert config.CONV == "flex_gemm" + assert config.ATTN == "flex_gemm_sparse_attn" + + from trellis2.backends import probe_metal_backends + + probes = probe_metal_backends() + assert probes["rasterizer"]["ok"], probes["rasterizer"] + assert probes["mesh_bvh"]["ok"], probes["mesh_bvh"] + + +@pytest.mark.skipif(platform.system() != "Darwin", reason="Metal parity requires macOS") +def test_pure_pytorch_sparse_conv_matches_flex_gemm_on_mps(): + code = r""" +import torch +from trellis2.modules.sparse import SparseTensor, config +from trellis2.modules.sparse.conv.conv import SparseConv3d + +if not torch.backends.mps.is_available(): + raise SystemExit(77) +coords = torch.tensor( + [[0, x, y, z] for x in range(3) for y in range(3) for z in range(3)], + dtype=torch.int32, + device='mps', +) +features = torch.randn(27, 4, dtype=torch.float16, device='mps') +sparse = SparseTensor( + feats=features, + coords=coords, + shape=torch.Size([1, 4]), + spatial_shape=[3, 3, 3], +) +config.set_conv_backend('flex_gemm') +metal = SparseConv3d(4, 5, 3, bias=True).half().to('mps') +metal_out = metal(sparse).feats + +config.set_conv_backend('pytorch') +fallback = SparseConv3d(4, 5, 3, bias=True).half().to('mps') +fallback.weight.data.copy_(metal.weight.data) +fallback.bias.data.copy_(metal.bias.data) +fallback_out = fallback(sparse).feats +torch.mps.synchronize() +print('PARITY=' + str(float((metal_out - fallback_out).abs().max().item()))) +""" + completed = subprocess.run( + [sys.executable, "-c", code], + cwd=ROOT, + text=True, + capture_output=True, + check=False, + ) + if completed.returncode == 77: + pytest.skip("MPS is unavailable") + assert completed.returncode == 0, completed.stderr + parity_line = next(line for line in completed.stdout.splitlines() if line.startswith("PARITY=")) + assert float(parity_line.removeprefix("PARITY=")) <= 2e-2 diff --git a/tests/test_background_preprocess.py b/tests/test_background_preprocess.py new file mode 100644 index 0000000..46a23fb --- /dev/null +++ b/tests/test_background_preprocess.py @@ -0,0 +1,72 @@ +import numpy as np +import pytest +from PIL import Image + +from trellis2.pipelines.trellis2_image_to_3d import Trellis2ImageTo3DPipeline + + +class _FakeRembg: + def __init__(self): + self.calls = 0 + + def __call__(self, image): + self.calls += 1 + rgba = image.convert("RGBA") + alpha = np.zeros((image.height, image.width), dtype=np.uint8) + alpha[2:6, 2:6] = 255 + rgba.putalpha(Image.fromarray(alpha, mode="L")) + return rgba + + +def _pipeline(rembg): + pipeline = object.__new__(Trellis2ImageTo3DPipeline) + pipeline.low_vram = False + pipeline.rembg_model = rembg + pipeline._device = "cpu" + return pipeline + + +def test_rgba_with_transparency_bypasses_rmbg(): + rembg = _FakeRembg() + rgba = np.full((8, 8, 4), 255, dtype=np.uint8) + rgba[:, :, :3] = [80, 120, 160] + rgba[:2, :, 3] = 0 + + output = _pipeline(rembg).preprocess_image(Image.fromarray(rgba, mode="RGBA")) + + assert rembg.calls == 0 + assert output.mode == "RGB" + assert output.size[0] > 0 and output.size[1] > 0 + + +def test_opaque_rgba_uses_official_background_remover(): + rembg = _FakeRembg() + opaque = Image.new("RGBA", (8, 8), (80, 120, 160, 255)) + + output = _pipeline(rembg).preprocess_image(opaque) + + assert rembg.calls == 1 + assert output.mode == "RGB" + assert output.size[0] > 0 and output.size[1] > 0 + + +def test_soft_authored_alpha_is_preserved_without_rmbg(): + rembg = _FakeRembg() + rgba = np.zeros((8, 8, 4), dtype=np.uint8) + rgba[2:6, 2:6, :3] = [80, 120, 160] + rgba[2:6, 2:6, 3] = 128 + + output = _pipeline(rembg).preprocess_image(Image.fromarray(rgba, mode="RGBA")) + + assert rembg.calls == 0 + assert output.size[0] > 0 and output.size[1] > 0 + + +def test_empty_authored_alpha_has_actionable_error(): + rembg = _FakeRembg() + empty = Image.new("RGBA", (8, 8), (80, 120, 160, 0)) + + with pytest.raises(RuntimeError, match="empty alpha mask"): + _pipeline(rembg).preprocess_image(empty) + + assert rembg.calls == 0 diff --git a/tests/test_generate_asset.py b/tests/test_generate_asset.py new file mode 100644 index 0000000..1968dad --- /dev/null +++ b/tests/test_generate_asset.py @@ -0,0 +1,128 @@ +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import torch + +from scripts import generate_asset + + +def _mesh(face_count): + return SimpleNamespace(faces=np.zeros((face_count, 3), dtype=np.int64)) + + +def test_default_has_no_decimation_limit(): + assert generate_asset._parse_target("none") is None + assert generate_asset._parse_target("0") is None + assert generate_asset._parse_target("250000") == 250000 + + +def test_oom_diagnostic_recommends_safe_single_process_retry(): + message = generate_asset._watchdog_message(RuntimeError("MPS backend out of memory")) + assert "--pipeline-type 512" in message + assert "one TRELLIS process" in message + + +def test_auto_backend_selects_mps_on_apple_silicon(monkeypatch): + monkeypatch.setattr(generate_asset.platform, "system", lambda: "Darwin") + monkeypatch.setattr(generate_asset.platform, "machine", lambda: "arm64") + assert generate_asset._resolve_backend("auto") == "mps" + + +def test_auto_backend_preserves_linux_cuda_route(monkeypatch): + monkeypatch.setattr(generate_asset.platform, "system", lambda: "Linux") + monkeypatch.setattr(generate_asset.platform, "machine", lambda: "x86_64") + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + assert generate_asset._resolve_backend("auto") == "cuda" + + +def test_auto_fallback_order_keeps_full_resolution_first(): + assert list(generate_asset._attempt_schedule("auto", None, 800_000)) == [ + ("metal", None), + ("kdtree", None), + ("metal", 200_000), + ("kdtree", 200_000), + ] + + +def test_simulated_bvh_failure_uses_full_kdtree(tmp_path: Path): + calls = [] + + def fake_export(mesh, path, *, baker, target, texture_size): + calls.append((baker, target, texture_size)) + if baker == "metal": + raise RuntimeError("simulated Metal BVH failure") + path.write_bytes(b"candidate") + return object(), False + + _, chosen, attempts = generate_asset._run_pbr_attempts( + _mesh(800_000), + tmp_path / "candidate_pbr.glb", + preferred_baker="auto", + requested_target=None, + texture_size=1024, + export_fn=fake_export, + ) + + assert calls == [("metal", None, 1024), ("kdtree", None, 1024)] + assert chosen["baker"] == "kdtree" + assert chosen["target_faces"] is None + assert [attempt["status"] for attempt in attempts] == ["failed", "ok"] + + +def test_safety_candidate_is_only_after_both_full_attempts(tmp_path: Path): + calls = [] + + def fake_export(mesh, path, *, baker, target, texture_size): + calls.append((baker, target)) + if target is None: + raise RuntimeError("full-resolution baker failure") + path.write_bytes(b"candidate") + return object(), True + + _, chosen, _ = generate_asset._run_pbr_attempts( + _mesh(800_000), + tmp_path / "candidate_pbr.glb", + preferred_baker="auto", + requested_target=None, + texture_size=512, + export_fn=fake_export, + ) + + assert calls == [("metal", None), ("kdtree", None), ("metal", 200_000)] + assert chosen["technical_safety_target"] is True + assert chosen["pre_simplified_before_bvh"] is True + + +def test_explicit_kdtree_never_loads_metal(): + assert list(generate_asset._attempt_schedule("kdtree", None, 1000)) == [("kdtree", None)] + + +def test_raw_export_preserves_full_topology_and_writes_vertex_normals(tmp_path: Path): + vertices = torch.tensor( + [ + [-1.0, -1.0, -1.0], + [1.0, -1.0, -1.0], + [1.0, 1.0, -1.0], + [-1.0, 1.0, -1.0], + [-1.0, -1.0, 1.0], + [1.0, -1.0, 1.0], + [1.0, 1.0, 1.0], + [-1.0, 1.0, 1.0], + ] + ) + faces = torch.tensor( + [ + [0, 1, 2], [0, 2, 3], [4, 6, 5], [4, 7, 6], + [0, 4, 5], [0, 5, 1], [1, 5, 6], [1, 6, 2], + [2, 6, 7], [2, 7, 3], [3, 7, 4], [3, 4, 0], + ] + ) + output = tmp_path / "raw_full.glb" + stats = generate_asset._export_raw(SimpleNamespace(vertices=vertices, faces=faces), output) + + from trellis2.gltf_validation import inspect_glb + + validated = inspect_glb(output, require_pbr=False) + assert stats["vertices"] == validated["vertices"] == 8 + assert stats["triangles"] == validated["triangles"] == 12 diff --git a/tests/test_mlx_parity.py b/tests/test_mlx_parity.py new file mode 100644 index 0000000..5038db6 --- /dev/null +++ b/tests/test_mlx_parity.py @@ -0,0 +1,103 @@ +import numpy as np +import pytest + +mx = pytest.importorskip("mlx.core") +torch = pytest.importorskip("torch") + +from mlx_backend.attention import MlxMultiHeadAttention +from mlx_backend.norm import LayerNorm32 +from mlx_backend.rope import apply_rope +from mlx_backend.sparse_conv import MlxSparseConv3d +from mlx_backend.sparse_tensor import MlxSparseTensor + + +def _numpy(value): + mx.eval(value) + return np.asarray(value) + + +def test_layer_norm_matches_torch_on_small_float32_input(): + rng = np.random.default_rng(42) + values = rng.normal(size=(3, 4, 8)).astype(np.float32) + weight = rng.normal(size=(8,)).astype(np.float32) + bias = rng.normal(size=(8,)).astype(np.float32) + + layer = LayerNorm32(8, eps=1e-6) + layer.weight = mx.array(weight) + layer.bias = mx.array(bias) + actual = _numpy(layer(mx.array(values))) + expected = torch.nn.functional.layer_norm( + torch.from_numpy(values), + (8,), + torch.from_numpy(weight), + torch.from_numpy(bias), + eps=1e-6, + ).numpy() + np.testing.assert_allclose(actual, expected, rtol=2e-5, atol=2e-5) + + +def test_rope_matches_complex_pair_reference(): + rng = np.random.default_rng(7) + values = rng.normal(size=(5, 2, 8)).astype(np.float32) + phases = rng.normal(size=(5, 4)).astype(np.float32) + cos = np.cos(phases) + sin = np.sin(phases) + + actual = _numpy(apply_rope(mx.array(values), mx.array(cos), mx.array(sin))) + paired = values.reshape(5, 2, 4, 2) + expected = np.empty_like(paired) + expected[..., 0] = paired[..., 0] * cos[:, None] - paired[..., 1] * sin[:, None] + expected[..., 1] = paired[..., 0] * sin[:, None] + paired[..., 1] * cos[:, None] + np.testing.assert_allclose(actual, expected.reshape(values.shape), rtol=1e-6, atol=1e-6) + + +def test_mlx_sdpa_matches_torch(): + rng = np.random.default_rng(11) + q = rng.normal(size=(6, 2, 4)).astype(np.float32) + k = rng.normal(size=(6, 2, 4)).astype(np.float32) + v = rng.normal(size=(6, 2, 4)).astype(np.float32) + attention = MlxMultiHeadAttention(channels=8, num_heads=2) + + actual = _numpy(attention._sdpa(mx.array(q), mx.array(k), mx.array(v))) + expected = torch.nn.functional.scaled_dot_product_attention( + torch.from_numpy(q).transpose(0, 1).unsqueeze(0), + torch.from_numpy(k).transpose(0, 1).unsqueeze(0), + torch.from_numpy(v).transpose(0, 1).unsqueeze(0), + scale=0.5, + )[0].transpose(0, 1).numpy() + np.testing.assert_allclose(actual, expected, rtol=3e-5, atol=3e-5) + + +def test_sparse_conv_matches_numpy_reference(): + rng = np.random.default_rng(19) + coords = np.array( + [ + [0, 0, 0, 0], + [0, 0, 0, 1], + [0, 0, 1, 0], + [0, 1, 0, 0], + [0, 1, 1, 1], + ], + dtype=np.int32, + ) + feats = rng.normal(size=(len(coords), 2)).astype(np.float32) + weight = rng.normal(size=(3, 3, 3, 3, 2)).astype(np.float32) + bias = rng.normal(size=(3,)).astype(np.float32) + + layer = MlxSparseConv3d(2, 3, kernel_size=3) + layer.weight = mx.array(weight) + layer.bias = mx.array(bias) + sparse = MlxSparseTensor(mx.array(feats), mx.array(coords), shape=(1, 2)) + actual = _numpy(layer(sparse).feats) + + lookup = {tuple(coord): index for index, coord in enumerate(coords)} + expected = np.tile(bias, (len(coords), 1)) + offsets = [(x, y, z) for x in (-1, 0, 1) for y in (-1, 0, 1) for z in (-1, 0, 1)] + flat_weight = weight.reshape(3, 27, 2).transpose(1, 2, 0) + for output_index, coord in enumerate(coords): + for kernel_index, offset in enumerate(offsets): + neighbor = (coord[0], coord[1] + offset[0], coord[2] + offset[1], coord[3] + offset[2]) + if neighbor in lookup: + expected[output_index] += feats[lookup[neighbor]] @ flat_weight[kernel_index] + + np.testing.assert_allclose(actual, expected, rtol=2e-5, atol=2e-5) diff --git a/tests/test_model_revisions.py b/tests/test_model_revisions.py new file mode 100644 index 0000000..ace6711 --- /dev/null +++ b/tests/test_model_revisions.py @@ -0,0 +1,36 @@ +from pathlib import Path + +from trellis2.model_revisions import ( + DINOV3_REPO, + DINOV3_REVISION, + MODEL_REVISIONS, + RMBG_REPO, + RMBG_REVISION, + SOURCE_REVISIONS, + TRELLIS_REPO, + TRELLIS_REVISION, + revision_for_repo, +) + +ROOT = Path(__file__).resolve().parents[1] + + +def test_runtime_revisions_are_full_commits(): + assert MODEL_REVISIONS[TRELLIS_REPO] == TRELLIS_REVISION + assert MODEL_REVISIONS[DINOV3_REPO] == DINOV3_REVISION + assert MODEL_REVISIONS[RMBG_REPO] == RMBG_REVISION + assert all(len(revision) == 40 for revision in MODEL_REVISIONS.values()) + assert all(len(revision) == 40 for revision in SOURCE_REVISIONS.values()) + + +def test_unknown_repo_keeps_explicit_revision(): + assert revision_for_repo("example/private-model", "release-1") == "release-1" + + +def test_source_revisions_are_present_in_reproducible_inputs(): + reproducibility_text = "\n".join( + (ROOT / path).read_text(encoding="utf-8") + for path in ("scripts/setup_macos.sh", "requirements_macos_core.txt", "README.md") + ) + for revision in SOURCE_REVISIONS.values(): + assert revision in reproducibility_text diff --git a/trellis2/backends.py b/trellis2/backends.py index 6ccae2e..2916817 100644 --- a/trellis2/backends.py +++ b/trellis2/backends.py @@ -8,6 +8,8 @@ Priority: Metal (mtldiffrast) > CUDA (nvdiffrast) > CPU fallback. Mesh: cumesh (auto-selects Metal/CUDA internally). """ import platform +import os +from functools import lru_cache import torch # --------------------------------------------------------------------------- @@ -15,6 +17,8 @@ import torch # --------------------------------------------------------------------------- HAS_MPS = hasattr(torch.backends, 'mps') and torch.backends.mps.is_available() HAS_CUDA = torch.cuda.is_available() +METAL_DISABLED = os.environ.get('TRELLIS_DISABLE_METAL', '0') == '1' +BACKEND_ERRORS = {} # --------------------------------------------------------------------------- # Differentiable rasterizer (dr) @@ -23,20 +27,21 @@ HAS_CUDA = torch.cuda.is_available() dr = None _dr_backend = None -try: - import mtldiffrast.torch as _mtldr - dr = _mtldr - _dr_backend = 'metal' -except ImportError: - pass +if not METAL_DISABLED: + try: + import mtldiffrast.torch as _mtldr + dr = _mtldr + _dr_backend = 'metal' + except (ImportError, RuntimeError, OSError) as exc: + BACKEND_ERRORS['mtldiffrast'] = str(exc) if dr is None: try: import nvdiffrast.torch as _nvdr dr = _nvdr _dr_backend = 'cuda' - except ImportError: - pass + except (ImportError, RuntimeError, OSError) as exc: + BACKEND_ERRORS['nvdiffrast'] = str(exc) HAS_DR = dr is not None @@ -58,14 +63,15 @@ 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 +if not METAL_DISABLED or platform.system() != 'Darwin': + 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, RuntimeError, OSError) as exc: + BACKEND_ERRORS['cumesh'] = str(exc) HAS_MESH = MeshBackend is not None HAS_REMESH = remesh_narrow_band_dc is not None @@ -93,15 +99,124 @@ except ImportError: _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 +if not METAL_DISABLED or platform.system() != 'Darwin': + try: + from flex_gemm.ops.grid_sample import grid_sample_3d as _fgs3d + _flex_grid_sample_3d = _fgs3d + HAS_FLEX_GEMM = True + except (ImportError, RuntimeError, OSError) as exc: + BACKEND_ERRORS['flex_gemm'] = str(exc) # --------------------------------------------------------------------------- # Overall backend name # --------------------------------------------------------------------------- BACKEND = _dr_backend or _mesh_backend or ('cpu' if HAS_TRIMESH else None) HAS_GPU = HAS_MPS or HAS_CUDA + + +def _probe_result(ok, error=None, **details): + return {'ok': bool(ok), 'error': error, **details} + + +@lru_cache(maxsize=1) +def probe_metal_rasterizer(): + """Run a tiny UV raster instead of treating a successful import as proof.""" + if METAL_DISABLED or _dr_backend != 'metal': + return _probe_result(False, 'Metal rasterizer is disabled or unavailable') + try: + positions = 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, + ) + triangles = torch.tensor([[0, 1, 2], [0, 2, 3]], dtype=torch.int32) + raster, _ = dr.rasterize( + RasterizeContext(), + positions, + triangles, + resolution=[8, 8], + ) + covered = int((raster[..., 3] != 0).sum().item()) + ok = tuple(raster.shape) == (1, 8, 8, 4) and covered == 64 + return _probe_result(ok, None if ok else 'unexpected raster output', covered_pixels=covered) + except Exception as exc: + return _probe_result(False, f'{type(exc).__name__}: {exc}') + + +@lru_cache(maxsize=1) +def probe_metal_mesh_bvh(): + """Exercise both MtlMesh and an actual nearest-point MtlBVH query.""" + if METAL_DISABLED or _mesh_backend != 'metal' or not HAS_MPS: + return _probe_result(False, 'Metal mesh/BVH backend is disabled or unavailable') + try: + vertices = torch.tensor( + [ + [-1.0, -1.0, -1.0], [1.0, -1.0, -1.0], + [1.0, 1.0, -1.0], [-1.0, 1.0, -1.0], + [-1.0, -1.0, 1.0], [1.0, -1.0, 1.0], + [1.0, 1.0, 1.0], [-1.0, 1.0, 1.0], + ], + dtype=torch.float32, + device='mps', + ) + faces = torch.tensor( + [ + [0, 1, 2], [0, 2, 3], [4, 6, 5], [4, 7, 6], + [0, 4, 5], [0, 5, 1], [1, 5, 6], [1, 6, 2], + [2, 6, 7], [2, 7, 3], [3, 7, 4], [3, 4, 0], + ], + dtype=torch.int32, + device='mps', + ) + mesh = MeshBackend() + mesh.init(vertices, faces) + mesh_vertices, mesh_faces = mesh.read() + bvh = BVH(vertices, faces) + query = torch.tensor([[0.0, 0.0, 1.2]], dtype=torch.float32, device='mps') + distance, face_id, uvw = bvh.unsigned_distance(query, return_uvw=True) + torch.mps.synchronize() + ok = ( + tuple(mesh_vertices.shape) == (8, 3) + and tuple(mesh_faces.shape) == (12, 3) + and bool(torch.isfinite(distance).all().item()) + and bool(torch.isfinite(uvw).all().item()) + and int(face_id.item()) >= 0 + ) + return _probe_result( + ok, + None if ok else 'unexpected mesh/BVH output', + distance=float(distance.cpu().item()), + face_id=int(face_id.cpu().item()), + ) + except Exception as exc: + return _probe_result(False, f'{type(exc).__name__}: {exc}') + + +def probe_metal_backends(): + return { + 'rasterizer': probe_metal_rasterizer(), + 'mesh_bvh': probe_metal_mesh_bvh(), + } + + +def backend_report(run_probes=True): + """Return a JSON-serializable snapshot used by setup and run metadata.""" + report = { + 'platform': platform.system(), + 'mps_available': HAS_MPS, + 'cuda_available': HAS_CUDA, + 'metal_disabled': METAL_DISABLED, + 'rasterizer': _dr_backend, + 'mesh': _mesh_backend, + 'flex_gemm': HAS_FLEX_GEMM, + 'trimesh': HAS_TRIMESH, + 'fast_simplification': HAS_FAST_SIMPLIFICATION, + 'errors': dict(BACKEND_ERRORS), + } + if run_probes and platform.system() == 'Darwin': + report['capability_probes'] = probe_metal_backends() + return report diff --git a/trellis2/gltf_validation.py b/trellis2/gltf_validation.py new file mode 100644 index 0000000..22d7fda --- /dev/null +++ b/trellis2/gltf_validation.py @@ -0,0 +1,151 @@ +"""Small glTF 2.0 validator for generated TRELLIS runtime outputs.""" + +from __future__ import annotations + +import base64 +import hashlib +import io +import json +import struct +from pathlib import Path +from typing import Any + +from PIL import Image + +JSON_CHUNK = 0x4E4F534A +BIN_CHUNK = 0x004E4942 + + +def _load_glb(path: Path) -> tuple[dict[str, Any], bytes, bytes]: + payload = path.read_bytes() + if len(payload) < 20: + raise ValueError(f"truncated GLB: {path}") + magic, version, declared_length = struct.unpack_from("<4sII", payload, 0) + if magic != b"glTF" or version != 2 or declared_length != len(payload): + raise ValueError(f"not a complete glTF 2.0 binary: {path}") + document = None + binary = b"" + offset = 12 + while offset < len(payload): + chunk_length, chunk_type = struct.unpack_from(" bytes: + view = document["bufferViews"][index] + start = int(view.get("byteOffset", 0)) + return binary[start : start + int(view["byteLength"])] + + +def _image(document: dict[str, Any], binary: bytes, index: int) -> bytes: + descriptor = document["images"][index] + if "bufferView" in descriptor: + return _view(document, binary, int(descriptor["bufferView"])) + uri = descriptor.get("uri", "") + if uri.startswith("data:"): + return base64.b64decode(uri.split(",", 1)[1]) + raise ValueError("all runtime textures must be embedded") + + +def _texture_payload(document: dict[str, Any], binary: bytes, info: dict[str, Any]) -> bytes: + texture = document["textures"][int(info["index"])] + return _image(document, binary, int(texture["source"])) + + +def inspect_glb(path: Path, *, require_pbr: bool) -> dict[str, Any]: + document, binary, payload = _load_glb(path) + if document.get("asset", {}).get("version") != "2.0": + raise ValueError("asset.version must be 2.0") + if document.get("animations") or document.get("cameras"): + raise ValueError("static outputs must not contain animations or cameras") + if document.get("extensions", {}).get("KHR_lights_punctual", {}).get("lights"): + raise ValueError("static outputs must not contain lights") + + vertices = 0 + triangles = 0 + bounds_min = [float("inf")] * 3 + bounds_max = [float("-inf")] * 3 + base_hashes = [] + mr_hashes = [] + primitive_count = 0 + for mesh in document.get("meshes", []): + for primitive in mesh.get("primitives", []): + primitive_count += 1 + if int(primitive.get("mode", 4)) != 4: + raise ValueError("only TRIANGLES primitives are supported") + attributes = primitive.get("attributes", {}) + required = ["POSITION", "NORMAL"] + (["TEXCOORD_0"] if require_pbr else []) + missing = [name for name in required if name not in attributes] + if missing: + raise ValueError(f"primitive is missing attributes: {', '.join(missing)}") + position = document["accessors"][int(attributes["POSITION"])] + if not position.get("min") or not position.get("max"): + raise ValueError("POSITION accessor has no bounds") + vertices += int(position["count"]) + for axis in range(3): + bounds_min[axis] = min(bounds_min[axis], float(position["min"][axis])) + bounds_max[axis] = max(bounds_max[axis], float(position["max"][axis])) + index_count = ( + int(document["accessors"][int(primitive["indices"])]["count"]) + if "indices" in primitive + else int(position["count"]) + ) + if index_count <= 0 or index_count % 3: + raise ValueError("invalid triangle index count") + triangles += index_count // 3 + + if require_pbr: + if "material" not in primitive: + raise ValueError("PBR primitive has no material") + material = document["materials"][int(primitive["material"])] + pbr = material.get("pbrMetallicRoughness", {}) + for key, hashes in ( + ("baseColorTexture", base_hashes), + ("metallicRoughnessTexture", mr_hashes), + ): + if key not in pbr: + raise ValueError(f"PBR material is missing {key}") + texture = _texture_payload(document, binary, pbr[key]) + hashes.append(hashlib.sha256(texture).hexdigest()) + with Image.open(io.BytesIO(_texture_payload(document, binary, pbr["baseColorTexture"]))) as image: + if "A" not in image.getbands(): + raise ValueError("base-color texture is missing generated alpha") + + if primitive_count == 0 or vertices == 0 or triangles == 0: + raise ValueError("GLB contains no non-empty mesh") + dimensions = [bounds_max[index] - bounds_min[index] for index in range(3)] + if any(dimension <= 0 for dimension in dimensions): + raise ValueError("GLB has empty bounds") + return { + "sha256": hashlib.sha256(payload).hexdigest(), + "bytes": len(payload), + "vertices": vertices, + "triangles": triangles, + "primitives": primitive_count, + "bounds_min": bounds_min, + "bounds_max": bounds_max, + "dimensions": dimensions, + "base_color_image_hashes": base_hashes, + "metallic_roughness_image_hashes": mr_hashes, + } + + +def validate_output_pair(raw_path: Path, candidate_path: Path) -> dict[str, Any]: + raw = inspect_glb(raw_path, require_pbr=False) + candidate = inspect_glb(candidate_path, require_pbr=True) + for raw_size, candidate_size in zip(raw["dimensions"], candidate["dimensions"]): + relative_delta = abs(raw_size - candidate_size) / max(raw_size, 1e-8) + if relative_delta > 0.05: + raise ValueError( + "raw_full.glb and candidate_pbr.glb bounds differ by more than 5%" + ) + return {"raw_full": raw, "candidate_pbr": candidate, "bounds_consistent": True} diff --git a/trellis2/model_revisions.py b/trellis2/model_revisions.py new file mode 100644 index 0000000..333023b --- /dev/null +++ b/trellis2/model_revisions.py @@ -0,0 +1,34 @@ +"""Pinned Hugging Face inputs used by the supported inference CLIs.""" + +TRELLIS_REPO = "microsoft/TRELLIS.2-4B" +TRELLIS_REVISION = "af44b45f2e35a493886929c6d786e563ec68364d" + +DINOV3_REPO = "facebook/dinov3-vitl16-pretrain-lvd1689m" +DINOV3_REVISION = "ea8dc2863c51be0a264bab82070e3e8836b02d51" + +RMBG_REPO = "briaai/RMBG-2.0" +RMBG_REVISION = "5df4c9c76d8170882c34f6986e848ee07fd0ba43" + +TRELLIS_IMAGE_LARGE_REPO = "microsoft/TRELLIS-image-large" +TRELLIS_IMAGE_LARGE_REVISION = "25e0d31ffbebe4b5a97464dd851910efc3002d96" + +MODEL_REVISIONS = { + TRELLIS_REPO: TRELLIS_REVISION, + TRELLIS_IMAGE_LARGE_REPO: TRELLIS_IMAGE_LARGE_REVISION, + DINOV3_REPO: DINOV3_REVISION, + RMBG_REPO: RMBG_REVISION, +} + +SOURCE_REVISIONS = { + "pedronaugusto/mtlgemm": "867aec8234299a7fe1ede7f802c8debe5a939a82", + "pedronaugusto/mtldiffrast": "4668cd91cb6d27f5e264731f94a06841fbf7aab8", + "pedronaugusto/mtlbvh": "23f441c470ce1f537e1fd836f3ffb5b8245f7975", + "pedronaugusto/mtlmesh": "212079e55772cff3d648a21372392c37e0643f3b", + "EasternJournalist/utils3d": "9a4eb15e4021b67b12c460c7057d642626897ec8", + "pedronaugusto/trellis2-apple": "6055b868734af6e12769d229d90580e775fae9f0", + "shivampkumar/trellis-mac": "d58628f4f5b9c3de8274cb110074154f4b31cef2", +} + + +def revision_for_repo(repo_id, default=None): + return MODEL_REVISIONS.get(repo_id, default) diff --git a/trellis2/models/__init__.py b/trellis2/models/__init__.py index d4fed03..58651ed 100644 --- a/trellis2/models/__init__.py +++ b/trellis2/models/__init__.py @@ -1,4 +1,7 @@ import importlib +from typing import Optional + +from ..model_revisions import revision_for_repo __attributes = { # Sparse Structure @@ -35,7 +38,14 @@ def __getattr__(name): return globals()[name] -def from_pretrained(path: str, **kwargs): +def from_pretrained( + path: str, + *, + revision: Optional[str] = None, + cache_dir: Optional[str] = None, + local_files_only: bool = False, + **kwargs, +): """ Load a model from a pretrained checkpoint. @@ -57,8 +67,14 @@ def from_pretrained(path: str, **kwargs): path_parts = path.split('/') repo_id = f'{path_parts[0]}/{path_parts[1]}' model_name = '/'.join(path_parts[2:]) - config_file = hf_hub_download(repo_id, f"{model_name}.json") - model_file = hf_hub_download(repo_id, f"{model_name}.safetensors") + revision = revision_for_repo(repo_id, revision) + hub_kwargs = { + "revision": revision, + "cache_dir": cache_dir, + "local_files_only": local_files_only, + } + config_file = hf_hub_download(repo_id, f"{model_name}.json", **hub_kwargs) + model_file = hf_hub_download(repo_id, f"{model_name}.safetensors", **hub_kwargs) with open(config_file, 'r') as f: config = json.load(f) diff --git a/trellis2/modules/image_feature_extractor.py b/trellis2/modules/image_feature_extractor.py index 8a635ac..8cb8772 100644 --- a/trellis2/modules/image_feature_extractor.py +++ b/trellis2/modules/image_feature_extractor.py @@ -8,6 +8,7 @@ from torchvision import transforms from transformers import DINOv3ViTModel import numpy as np from PIL import Image +from ..model_revisions import DINOV3_REPO, DINOV3_REVISION class DinoV2FeatureExtractor: @@ -69,9 +70,21 @@ class DinoV3FeatureExtractor: """ Feature extractor for DINOv3 models. """ - def __init__(self, model_name: str, image_size=512): + def __init__( + self, + model_name: str = DINOV3_REPO, + image_size=512, + revision: Optional[str] = DINOV3_REVISION, + cache_dir: Optional[str] = None, + local_files_only: bool = False, + ): self.model_name = model_name - self.model = DINOv3ViTModel.from_pretrained(model_name) + self.model = DINOv3ViTModel.from_pretrained( + model_name, + revision=revision, + cache_dir=cache_dir, + local_files_only=local_files_only, + ) self.model.eval() self.image_size = image_size self.transform = transforms.Compose([ diff --git a/trellis2/modules/sparse/config.py b/trellis2/modules/sparse/config.py index 303f026..f091dbd 100644 --- a/trellis2/modules/sparse/config.py +++ b/trellis2/modules/sparse/config.py @@ -1,4 +1,5 @@ from typing import * +import math import platform import sys @@ -12,12 +13,15 @@ def __detect_defaults(): if platform.system() == 'Darwin': if __flex_gemm_works_on_mps(): CONV = 'flex_gemm' - # flex_gemm_sparse_attn is now flash-attention-v2 with simdgroup - # matmul + simd-shuffle softmax (5–15× over SDPA-padded-CPU-bounce - # at every measured shape including max_seqlen=2048). The conv - # probe above covers the same package, so if it passed, the - # attention path is available too. - ATTN = 'flex_gemm_sparse_attn' + # Sparse convolution and sparse attention are separate Metal + # kernels. A working convolution install does not prove that the + # attention entry point is ABI-compatible with the active torch + # build, so select it only after its own numerical probe. + ATTN = ( + 'flex_gemm_sparse_attn' + if probe_flex_gemm_sparse_attention_on_mps() + else 'sdpa' + ) else: CONV = 'pytorch' ATTN = 'sdpa' @@ -35,6 +39,9 @@ def __flex_gemm_works_on_mps(): and move to MPS because some PyTorch builds lack int/fp16 torch.zeros kernels on MPS.""" try: + import os + if os.environ.get('TRELLIS_DISABLE_METAL', '0') == '1': + return False import torch if not torch.backends.mps.is_available(): return False @@ -58,6 +65,60 @@ def __flex_gemm_works_on_mps(): except Exception: return False + +def probe_flex_gemm_sparse_attention_on_mps() -> bool: + """Exercise the real Metal sparse-attention kernel and compare it to SDPA. + + This intentionally uses the production head dimension (64) while keeping + the sequence tiny. Returning ``False`` is a controlled capability result: + callers fall back to PyTorch SDPA without disabling working Metal sparse + convolution kernels. + """ + try: + import os + if os.environ.get('TRELLIS_DISABLE_METAL', '0') == '1': + return False + + import torch + import torch.nn.functional as F + if not torch.backends.mps.is_available(): + return False + + import flex_gemm + + tokens, heads, head_dim = 16, 2, 64 + generator = torch.Generator(device='cpu').manual_seed(42) + q = torch.randn(tokens, heads, head_dim, dtype=torch.float16, generator=generator).to('mps').contiguous() + k = torch.randn(tokens, heads, head_dim, dtype=torch.float16, generator=generator).to('mps').contiguous() + v = torch.randn(tokens, heads, head_dim, dtype=torch.float16, generator=generator).to('mps').contiguous() + cu_seqlens = torch.tensor([0, tokens], dtype=torch.int32).to('mps') + + out = flex_gemm.kernels.cuda.sparse_attention_fwd( + q, + k, + v, + cu_seqlens, + cu_seqlens, + tokens, + tokens, + 1.0 / math.sqrt(head_dim), + ) + reference = F.scaled_dot_product_attention( + q.transpose(0, 1).unsqueeze(0), + k.transpose(0, 1).unsqueeze(0), + v.transpose(0, 1).unsqueeze(0), + ).squeeze(0).transpose(0, 1) + torch.mps.synchronize() + + return ( + out.device.type == 'mps' + and out.shape == reference.shape + and bool(torch.isfinite(out).all().item()) + and bool(torch.allclose(out, reference, rtol=2e-2, atol=2e-2)) + ) + except Exception: + return False + def __has_cuda(): try: import torch diff --git a/trellis2/pipelines/__init__.py b/trellis2/pipelines/__init__.py index 9bb6835..0166945 100644 --- a/trellis2/pipelines/__init__.py +++ b/trellis2/pipelines/__init__.py @@ -1,4 +1,5 @@ import importlib +from ..model_revisions import revision_for_repo __attributes = { "Trellis2ImageTo3DPipeline": "trellis2_image_to_3d", @@ -23,7 +24,7 @@ def __getattr__(name): return globals()[name] -def from_pretrained(path: str): +def from_pretrained(path: str, **hub_kwargs): """ Load a pipeline from a model folder or a Hugging Face model hub. @@ -38,11 +39,13 @@ def from_pretrained(path: str): config_file = f"{path}/pipeline.json" else: from huggingface_hub import hf_hub_download - config_file = hf_hub_download(path, "pipeline.json") + if hub_kwargs.get("revision") is None: + hub_kwargs["revision"] = revision_for_repo(path) + config_file = hf_hub_download(path, "pipeline.json", **hub_kwargs) with open(config_file, 'r') as f: config = json.load(f) - return globals()[config['name']].from_pretrained(path) + return globals()[config['name']].from_pretrained(path, **hub_kwargs) # For PyLance diff --git a/trellis2/pipelines/base.py b/trellis2/pipelines/base.py index 331e1ed..cdb2c8e 100644 --- a/trellis2/pipelines/base.py +++ b/trellis2/pipelines/base.py @@ -2,6 +2,7 @@ from typing import * import torch import torch.nn as nn from .. import models +from ..model_revisions import revision_for_repo class Pipeline: @@ -19,7 +20,15 @@ class Pipeline: model.eval() @classmethod - def from_pretrained(cls, path: str, config_file: str = "pipeline.json") -> "Pipeline": + def from_pretrained( + cls, + path: str, + config_file: str = "pipeline.json", + *, + revision: Optional[str] = None, + cache_dir: Optional[str] = None, + local_files_only: bool = False, + ) -> "Pipeline": """ Load a pretrained model. """ @@ -31,7 +40,14 @@ class Pipeline: config_file = f"{path}/{config_file}" else: from huggingface_hub import hf_hub_download - config_file = hf_hub_download(path, config_file) + revision = revision_for_repo(path, revision) + config_file = hf_hub_download( + path, + config_file, + revision=revision, + cache_dir=cache_dir, + local_files_only=local_files_only, + ) with open(config_file, 'r') as f: args = json.load(f)['args'] @@ -40,13 +56,22 @@ class Pipeline: for k, v in args['models'].items(): if hasattr(cls, 'model_names_to_load') and k not in cls.model_names_to_load: continue - try: - _models[k] = models.from_pretrained(f"{path}/{v}") - except Exception as e: - _models[k] = models.from_pretrained(v) + is_external_model = v.count('/') >= 2 + model_path = v if is_external_model else f"{path}/{v}" + _models[k] = models.from_pretrained( + model_path, + revision=None if is_external_model else revision, + cache_dir=cache_dir, + local_files_only=local_files_only, + ) new_pipeline = cls(_models) new_pipeline._pretrained_args = args + new_pipeline._pretrained_hub_kwargs = { + "revision": revision, + "cache_dir": cache_dir, + "local_files_only": local_files_only, + } return new_pipeline @property @@ -69,4 +94,4 @@ class Pipeline: self.to(torch.device("cuda")) def cpu(self) -> None: - self.to(torch.device("cpu")) \ No newline at end of file + self.to(torch.device("cpu")) diff --git a/trellis2/pipelines/rembg/BiRefNet.py b/trellis2/pipelines/rembg/BiRefNet.py index 61056fc..effb057 100644 --- a/trellis2/pipelines/rembg/BiRefNet.py +++ b/trellis2/pipelines/rembg/BiRefNet.py @@ -3,12 +3,23 @@ from transformers import AutoModelForImageSegmentation import torch from torchvision import transforms from PIL import Image +from ...model_revisions import RMBG_REPO, RMBG_REVISION class BiRefNet: - def __init__(self, model_name: str = "ZhengPeng7/BiRefNet"): + def __init__( + self, + model_name: str = RMBG_REPO, + revision: Optional[str] = RMBG_REVISION, + cache_dir: Optional[str] = None, + local_files_only: bool = False, + ): self.model = AutoModelForImageSegmentation.from_pretrained( - model_name, trust_remote_code=True, + model_name, + revision=revision, + cache_dir=cache_dir, + local_files_only=local_files_only, + trust_remote_code=True, ) self.model.eval() self.transform_image = transforms.Compose( @@ -40,4 +51,3 @@ class BiRefNet: mask = pred_pil.resize(image_size) image.putalpha(mask) return image - \ No newline at end of file diff --git a/trellis2/pipelines/trellis2_image_to_3d.py b/trellis2/pipelines/trellis2_image_to_3d.py index 465986b..a6d25ef 100644 --- a/trellis2/pipelines/trellis2_image_to_3d.py +++ b/trellis2/pipelines/trellis2_image_to_3d.py @@ -89,15 +89,24 @@ class Trellis2ImageTo3DPipeline(Pipeline): self._device = 'cpu' @classmethod - def from_pretrained(cls, path: str, config_file: str = "pipeline.json") -> "Trellis2ImageTo3DPipeline": + def from_pretrained( + cls, + path: str, + config_file: str = "pipeline.json", + **hub_kwargs, + ) -> "Trellis2ImageTo3DPipeline": """ Load a pretrained model. Args: path (str): The path to the model. Can be either local path or a Hugging Face repository. """ - pipeline = super().from_pretrained(path, config_file) + pipeline = super().from_pretrained(path, config_file, **hub_kwargs) args = pipeline._pretrained_args + dependency_hub_kwargs = { + "cache_dir": hub_kwargs.get("cache_dir"), + "local_files_only": hub_kwargs.get("local_files_only", False), + } 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'] @@ -111,8 +120,14 @@ class Trellis2ImageTo3DPipeline(Pipeline): pipeline.shape_slat_normalization = args['shape_slat_normalization'] pipeline.tex_slat_normalization = args['tex_slat_normalization'] - pipeline.image_cond_model = getattr(image_feature_extractor, args['image_cond_model']['name'])(**args['image_cond_model']['args']) - pipeline.rembg_model = getattr(rembg, args['rembg_model']['name'])(**args['rembg_model']['args']) + pipeline.image_cond_model = getattr(image_feature_extractor, args['image_cond_model']['name'])( + **args['image_cond_model']['args'], + **dependency_hub_kwargs, + ) + pipeline.rembg_model = getattr(rembg, args['rembg_model']['name'])( + **args['rembg_model']['args'], + **dependency_hub_kwargs, + ) pipeline.low_vram = args.get('low_vram', True) pipeline.default_pipeline_type = args.get('default_pipeline_type', '1024_cascade') @@ -161,6 +176,16 @@ class Trellis2ImageTo3DPipeline(Pipeline): output_np = np.array(output) alpha = output_np[:, :, 3] bbox = np.argwhere(alpha > 0.8 * 255) + # Preserve deliberately soft RGBA masks. The high-confidence threshold + # is useful for RMBG output, but an authored semi-transparent object can + # legitimately have no pixels above it. + if bbox.size == 0: + bbox = np.argwhere(alpha > 0) + if bbox.size == 0: + raise RuntimeError( + "Background preprocessing produced an empty alpha mask; provide " + "a non-empty RGBA input or retry with --background keep." + ) bbox = np.min(bbox[:, 1]), np.min(bbox[:, 0]), np.max(bbox[:, 1]), np.max(bbox[:, 0]) center = (bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2 size = max(bbox[2] - bbox[0], bbox[3] - bbox[1]) diff --git a/trellis2/pipelines/trellis2_texturing.py b/trellis2/pipelines/trellis2_texturing.py index 4fa5d64..ecf6336 100755 --- a/trellis2/pipelines/trellis2_texturing.py +++ b/trellis2/pipelines/trellis2_texturing.py @@ -68,15 +68,24 @@ class Trellis2TexturingPipeline(Pipeline): self._device = 'cpu' @classmethod - def from_pretrained(cls, path: str, config_file: str = "pipeline.json") -> "Trellis2TexturingPipeline": + def from_pretrained( + cls, + path: str, + config_file: str = "pipeline.json", + **hub_kwargs, + ) -> "Trellis2TexturingPipeline": """ Load a pretrained model. Args: path (str): The path to the model. Can be either local path or a Hugging Face repository. """ - pipeline = super().from_pretrained(path, config_file) + pipeline = super().from_pretrained(path, config_file, **hub_kwargs) args = pipeline._pretrained_args + dependency_hub_kwargs = { + "cache_dir": hub_kwargs.get("cache_dir"), + "local_files_only": hub_kwargs.get("local_files_only", False), + } 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'] @@ -84,8 +93,14 @@ class Trellis2TexturingPipeline(Pipeline): pipeline.shape_slat_normalization = args['shape_slat_normalization'] pipeline.tex_slat_normalization = args['tex_slat_normalization'] - pipeline.image_cond_model = getattr(image_feature_extractor, args['image_cond_model']['name'])(**args['image_cond_model']['args']) - pipeline.rembg_model = getattr(rembg, args['rembg_model']['name'])(**args['rembg_model']['args']) + pipeline.image_cond_model = getattr(image_feature_extractor, args['image_cond_model']['name'])( + **args['image_cond_model']['args'], + **dependency_hub_kwargs, + ) + pipeline.rembg_model = getattr(rembg, args['rembg_model']['name'])( + **args['rembg_model']['args'], + **dependency_hub_kwargs, + ) pipeline.low_vram = args.get('low_vram', True) pipeline.pbr_attr_layout = { diff --git a/trellis2/representations/mesh/base.py b/trellis2/representations/mesh/base.py index 857efbf..6f8fc10 100644 --- a/trellis2/representations/mesh/base.py +++ b/trellis2/representations/mesh/base.py @@ -2,6 +2,7 @@ from typing import * import torch import torch.nn.functional as F_grid import numpy as np +import os from ..voxel import Voxel from trellis2.backends import ( @@ -46,7 +47,10 @@ class Mesh: return self.to('cpu') def _gpu_device(self): - if _HAS_MPS: + # Decoder-sized meshes have historically triggered watchdogs or + # process-level crashes in Metal cleanup kernels. Keep the supported + # default on the safe CPU implementations; advanced users may opt in. + if _HAS_MPS and os.environ.get('TRELLIS_ENABLE_MPS_DECODE_MESH_OPS') == '1': return 'mps' if _HAS_CUDA: return 'cuda'