223 lines
7.7 KiB
Python
223 lines
7.7 KiB
Python
"""
|
|
Centralized backend resolution for trellis2-apple.
|
|
|
|
Resolves differentiable rasterization (dr), mesh processing (MeshBackend, BVH),
|
|
remeshing, and grid_sample once at import time. All other modules import from here.
|
|
|
|
Priority: Metal (mtldiffrast) > CUDA (nvdiffrast) > CPU fallback.
|
|
Mesh: cumesh (auto-selects Metal/CUDA internally).
|
|
"""
|
|
import platform
|
|
import os
|
|
from functools import lru_cache
|
|
import torch
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Detect platform
|
|
# ---------------------------------------------------------------------------
|
|
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)
|
|
# mtldiffrast (Metal) and nvdiffrast (CUDA) are separate packages
|
|
# ---------------------------------------------------------------------------
|
|
dr = None
|
|
_dr_backend = None
|
|
|
|
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, RuntimeError, OSError) as exc:
|
|
BACKEND_ERRORS['nvdiffrast'] = str(exc)
|
|
|
|
HAS_DR = dr is not None
|
|
|
|
|
|
def RasterizeContext(device=None):
|
|
"""Create the appropriate rasterization context for the active backend."""
|
|
if _dr_backend == 'metal':
|
|
return dr.MtlRasterizeContext(device=device)
|
|
elif _dr_backend == 'cuda':
|
|
return dr.RasterizeCudaContext(device=device)
|
|
raise RuntimeError("No differentiable rasterization backend available")
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Mesh processing (MeshBackend, BVH, remesh)
|
|
# cumesh auto-selects Metal or CUDA backend internally
|
|
# ---------------------------------------------------------------------------
|
|
MeshBackend = None
|
|
BVH = None
|
|
remesh_narrow_band_dc = None
|
|
_mesh_backend = None
|
|
|
|
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
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# CPU fallback mesh libraries
|
|
# ---------------------------------------------------------------------------
|
|
try:
|
|
import trimesh
|
|
HAS_TRIMESH = True
|
|
except ImportError:
|
|
trimesh = None
|
|
HAS_TRIMESH = False
|
|
|
|
try:
|
|
import fast_simplification
|
|
HAS_FAST_SIMPLIFICATION = True
|
|
except ImportError:
|
|
fast_simplification = None
|
|
HAS_FAST_SIMPLIFICATION = False
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Grid sample (flex_gemm on Metal/CUDA, F.grid_sample fallback)
|
|
# ---------------------------------------------------------------------------
|
|
_flex_grid_sample_3d = None
|
|
HAS_FLEX_GEMM = False
|
|
|
|
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
|