feat: add native Apple Silicon asset pipeline
This commit is contained in:
Jourloy
parent
9c34570ffc
commit
f59e6e5981
94
README.md
94
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
|
||||
|
||||
|
||||
@ -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
|
||||
168
api_server.py
168
api_server.py
@ -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()
|
||||
174
app_mlx.py
174
app_mlx.py
@ -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()
|
||||
@ -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)
|
||||
|
||||
@ -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')
|
||||
|
||||
@ -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")
|
||||
|
||||
|
||||
@ -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")
|
||||
|
||||
|
||||
@ -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 <url>)
|
||||
mtldiffrast @ https://github.com/pedronaugusto/mtldiffrast/archive/main.tar.gz
|
||||
cumesh @ https://github.com/pedronaugusto/mtlmesh/archive/main.tar.gz
|
||||
flex_gemm @ https://github.com/pedronaugusto/mtlgemm/archive/main.tar.gz
|
||||
|
||||
# Misc
|
||||
tqdm
|
||||
imageio
|
||||
imageio-ffmpeg
|
||||
scipy
|
||||
einops
|
||||
# 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
|
||||
|
||||
33
requirements_macos_core.txt
Normal file
33
requirements_macos_core.txt
Normal file
@ -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
|
||||
118
scripts/download_weights.py
Normal file → Executable file
118
scripts/download_weights.py
Normal file → Executable file
@ -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())
|
||||
|
||||
574
scripts/generate_asset.py
Executable file
574
scripts/generate_asset.py
Executable file
@ -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())
|
||||
118
scripts/probe_macos.py
Executable file
118
scripts/probe_macos.py
Executable file
@ -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())
|
||||
118
scripts/setup_macos.sh
Executable file
118
scripts/setup_macos.sh
Executable file
@ -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
|
||||
28
scripts/validate_outputs.py
Executable file
28
scripts/validate_outputs.py
Executable file
@ -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())
|
||||
@ -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.")
|
||||
|
||||
6
tests/conftest.py
Normal file
6
tests/conftest.py
Normal file
@ -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"))
|
||||
114
tests/test_backend_fallbacks.py
Normal file
114
tests/test_backend_fallbacks.py
Normal file
@ -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
|
||||
72
tests/test_background_preprocess.py
Normal file
72
tests/test_background_preprocess.py
Normal file
@ -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
|
||||
128
tests/test_generate_asset.py
Normal file
128
tests/test_generate_asset.py
Normal file
@ -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
|
||||
103
tests/test_mlx_parity.py
Normal file
103
tests/test_mlx_parity.py
Normal file
@ -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)
|
||||
36
tests/test_model_revisions.py
Normal file
36
tests/test_model_revisions.py
Normal file
@ -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
|
||||
@ -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
|
||||
|
||||
151
trellis2/gltf_validation.py
Normal file
151
trellis2/gltf_validation.py
Normal file
@ -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("<II", payload, offset)
|
||||
offset += 8
|
||||
chunk = payload[offset : offset + chunk_length]
|
||||
offset += chunk_length
|
||||
if chunk_type == JSON_CHUNK:
|
||||
document = json.loads(chunk.rstrip(b" \t\r\n\0").decode("utf-8"))
|
||||
elif chunk_type == BIN_CHUNK:
|
||||
binary = chunk
|
||||
if document is None:
|
||||
raise ValueError(f"missing GLB JSON chunk: {path}")
|
||||
return document, binary, payload
|
||||
|
||||
|
||||
def _view(document: dict[str, Any], binary: bytes, index: int) -> 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}
|
||||
34
trellis2/model_revisions.py
Normal file
34
trellis2/model_revisions.py
Normal file
@ -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)
|
||||
@ -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)
|
||||
|
||||
@ -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([
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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])
|
||||
|
||||
@ -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 = {
|
||||
|
||||
@ -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'
|
||||
|
||||
Loading…
Reference in New Issue
Block a user