""" blackwell_fix.py — TRELLIS.2 Mesh Extraction Fix for NVIDIA Blackwell GPUs Drop-in workaround for the broken CuMesh remeshing pipeline on sm_120 GPUs (RTX 5070, 5070 Ti, 5080, 5090). Replaces the CUDA-dependent to_glb() mesh extraction with a voxel-based marching cubes approach that produces watertight, 3D-printable meshes using only CPU operations. Usage: import blackwell_fix # Auto-detect Blackwell and apply all compatibility patches blackwell_fix.patch_all() # After running TRELLIS.2 inference: mesh = pipeline.run(image)[0] trimesh_mesh = blackwell_fix.voxel_to_mesh(mesh) trimesh_mesh.export("output.stl") Requirements: numpy, scipy, scikit-image, trimesh """ import os import sys import gc import time from typing import Optional import numpy as np # --------------------------------------------------------------------------- # Blackwell GPU detection # --------------------------------------------------------------------------- def is_blackwell_gpu(device: int = 0) -> bool: """Check if the current GPU is a Blackwell-architecture device (sm_120).""" try: import torch if not torch.cuda.is_available(): return False major, minor = torch.cuda.get_device_capability(device) return major >= 12 except Exception: return False def get_gpu_info(device: int = 0) -> dict: """Return GPU name and compute capability.""" try: import torch if not torch.cuda.is_available(): return {"name": "N/A", "compute_capability": (0, 0), "is_blackwell": False} props = torch.cuda.get_device_properties(device) cc = torch.cuda.get_device_capability(device) return { "name": props.name, "compute_capability": cc, "is_blackwell": cc[0] >= 12, } except Exception: return {"name": "N/A", "compute_capability": (0, 0), "is_blackwell": False} # --------------------------------------------------------------------------- # Compatibility patches — must be applied BEFORE importing TRELLIS.2 # --------------------------------------------------------------------------- _patches_applied = False def patch_all(force: bool = False, verbose: bool = True): """ Apply all Blackwell compatibility patches. Must be called BEFORE importing trellis2 or o_voxel. Sets environment variables and monkey-patches CUDA capability detection so that spconv, cumm, and flex_gemm select sm_90 (Hopper) PTX kernels, which JIT-compile correctly on Blackwell hardware. Args: force: Apply patches even on non-Blackwell GPUs (for testing). verbose: Print status messages for each patch applied. """ global _patches_applied if _patches_applied: return _patches_applied = True import torch # Check if patches are needed major, minor = torch.cuda.get_device_capability(0) if torch.cuda.is_available() else (0, 0) if major < 10 and not force: if verbose: print(f"[blackwell_fix] GPU CC {major}.{minor} — no patches needed") return if verbose: name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "unknown" print(f"[blackwell_fix] Detected {name} (CC {major}.{minor}) — applying patches") # ── Environment variables (must be set before TRELLIS imports) ── os.environ["ATTN_BACKEND"] = "sdpa" # PyTorch native SDPA attention os.environ["SPARSE_CONV_BACKEND"] = "spconv" # Avoid Triton (broken on CC 12.0) os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" # ── 1) Patch torch.cuda.get_device_capability ── _orig_cap = torch.cuda.get_device_capability def _patched_cap(device=None): m, n = _orig_cap(device) return (9, 0) if m >= 10 else (m, n) torch.cuda.get_device_capability = _patched_cap # ── 2) Patch cumm/spconv compute capability detection ── try: import cumm.tensorview as _tv _orig_cumm = _tv.get_compute_capability def _patched_cumm(index: int = -1): m, n = _orig_cumm(index) return (9, 0) if m >= 10 else (m, n) _tv.get_compute_capability = _patched_cumm try: import cumm.tensorview_bind as _tvb _orig_tvb = _tvb.get_compute_capability def _patched_tvb(index: int = -1): m, n = _orig_tvb(index) return (9, 0) if m >= 10 else (m, n) _tvb.get_compute_capability = _patched_tvb except (ImportError, AttributeError): pass if verbose: print("[blackwell_fix] Patched cumm CC detection -> (9, 0)") except ImportError: pass # ── 3) Patch flex_gemm Triton kernels with PyTorch fallbacks ── # Triton 3.3.x cannot compile for CC >= 10.0 try: import flex_gemm.kernels.triton as _fgk def _fwd(feats, indices, weight): idx = indices.long().clamp(0, feats.shape[0] - 1) return (feats[idx] * weight.unsqueeze(-1)).sum(dim=1) def _bwd(grad_output, indices, weight, N): M, C = grad_output.shape idx = indices.long().clamp(0, N - 1) wg = grad_output.unsqueeze(1) * weight.unsqueeze(-1) gf = torch.zeros(N, C, device=grad_output.device, dtype=grad_output.dtype) gf.scatter_add_(0, idx.unsqueeze(-1).expand_as(wg).reshape(-1, C), wg.reshape(-1, C)) return gf _fgk.indice_weighed_sum_fwd = _fwd _fgk.indice_weighed_sum_bwd_input = _bwd if verbose: print("[blackwell_fix] Patched flex_gemm Triton kernels -> PyTorch") except (ImportError, AttributeError): pass # ── 4) Pillow WebP compatibility ── try: from PIL import _webp if not hasattr(_webp, 'HAVE_WEBPANIM'): _webp.HAVE_WEBPANIM = hasattr(_webp, 'WebPAnimDecoder') if not hasattr(_webp, 'HAVE_WEBPMUX'): _webp.HAVE_WEBPMUX = hasattr(_webp, 'WebPAnimDecoder') if not hasattr(_webp, 'HAVE_TRANSPARENCY'): _webp.HAVE_TRANSPARENCY = True except (ImportError, AttributeError): pass # --------------------------------------------------------------------------- # Voxel-based mesh extraction (replaces broken CuMesh pipeline) # --------------------------------------------------------------------------- def voxel_to_mesh( mesh_output, target_height_mm: float = 100.0, sigma: float = 1.5, coarse_downsample: float = 4, taubin_iterations: int = 50, verbose: bool = True, ): """ Convert a TRELLIS.2 mesh output to a watertight trimesh via voxel marching cubes. This replaces the broken to_glb() → CuMesh pipeline on Blackwell GPUs. Uses a two-phase approach: Phase 1: Coarse grid (downsampled) with aggressive morphological closing + flood fill to determine the solid interior. Phase 2: Full-resolution voxel grid combining surface voxels with the upscaled coarse interior. Gaussian smooth + marching cubes. Args: mesh_output: TRELLIS.2 mesh object (from pipeline.run()[0]). Must have .coords (integer voxel positions) and .voxel_size (float scaling factor). target_height_mm: Scale the output so the tallest dimension equals this many millimeters. Set to 0 to skip scaling. sigma: Gaussian smoothing sigma for the volume before marching cubes. Higher = smoother surface but less detail. 1.0-2.0 recommended. coarse_downsample: Downsample factor for the coarse interior fill. 4 works well for most models. taubin_iterations: Number of Taubin mesh smoothing passes after marching cubes. Reduces voxel staircase artifacts. verbose: Print progress messages. Returns: trimesh.Trimesh: Watertight mesh ready for export/printing. """ import torch # Extract numpy arrays from the TRELLIS mesh object if isinstance(mesh_output.coords, torch.Tensor): coords_np = mesh_output.coords.cpu().numpy().copy() else: coords_np = np.array(mesh_output.coords).copy() voxel_size = float(mesh_output.voxel_size) return voxel_coords_to_mesh( coords_np=coords_np, voxel_size=voxel_size, target_height_mm=target_height_mm, sigma=sigma, coarse_downsample=coarse_downsample, taubin_iterations=taubin_iterations, verbose=verbose, ) def voxel_coords_to_mesh( coords_np: np.ndarray, voxel_size: float, target_height_mm: float = 100.0, sigma: float = 1.5, coarse_downsample: float = 4, taubin_iterations: int = 50, verbose: bool = True, ): """ Convert raw voxel coordinates to a watertight trimesh. This is the lower-level version of voxel_to_mesh() that works directly with numpy arrays. Useful when you've already extracted and saved the coordinates (e.g., for iterating on mesh parameters without re-running inference). Args: coords_np: Integer voxel coordinates, shape (N, 3). voxel_size: Voxel size in model units. target_height_mm: Scale output height. Set to 0 to skip. sigma: Gaussian smoothing sigma. coarse_downsample: Coarse grid downsample factor. taubin_iterations: Mesh smoothing iterations. verbose: Print progress. Returns: trimesh.Trimesh: Watertight mesh. """ import trimesh from scipy import ndimage from scipy.ndimage import gaussian_filter, zoom from skimage import measure from scipy.sparse import coo_matrix from scipy.sparse.csgraph import connected_components def log(msg): if verbose: print(f" {msg}") t0 = time.time() log("Multi-resolution voxel reconstruction...") c_min = coords_np.min(axis=0) c_max = coords_np.max(axis=0) extent = (c_max - c_min).astype(int) log(f"Voxel extent: {extent[0]}x{extent[1]}x{extent[2]} " f"({coords_np.shape[0]:,} occupied, voxel_size={voxel_size:.6f})") struct26 = ndimage.generate_binary_structure(3, 3) # 26-connected ds_c = coarse_downsample # ── Phase 1: Coarse fill for solid interior ────────────────────────── c_ds = ((coords_np - c_min) // ds_c).astype(int) grid_c = tuple(((c_max - c_min) // ds_c + 2).astype(int)) vol_c = np.zeros(grid_c, dtype=np.uint8) vol_c[c_ds[:, 0], c_ds[:, 1], c_ds[:, 2]] = 1 del c_ds log(f"Phase 1 — Coarse grid ({ds_c}x): {grid_c[0]}x{grid_c[1]}x{grid_c[2]}") # Aggressive morphological closing to seal all gaps vol_c = ndimage.binary_dilation(vol_c, struct26, iterations=5) pad_c = 3 vol_c = np.pad(vol_c.astype(np.uint8), pad_c, mode='constant', constant_values=0) vol_c = ndimage.binary_fill_holes(vol_c) vol_c = ndimage.binary_erosion(vol_c, struct26, iterations=4) vol_c = vol_c[pad_c:-pad_c, pad_c:-pad_c, pad_c:-pad_c] log(f"Coarse interior: {int(vol_c.sum()):,} voxels") # ── Phase 2: Full-resolution grid + coarse interior ────────────────── pad_f = 2 grid_f = tuple((c_max - c_min + 2 + 2 * pad_f).astype(int)) log(f"Phase 2 — Fine grid (ds=1): {grid_f[0]}x{grid_f[1]}x{grid_f[2]} " f"({int(np.prod(grid_f))/1e6:.0f}M cells)") # Place surface voxels c_shifted = (coords_np - c_min + pad_f).astype(int) vol = np.zeros(grid_f, dtype=np.uint8) vol[c_shifted[:, 0], c_shifted[:, 1], c_shifted[:, 2]] = 1 del c_shifted, coords_np gc.collect() surface = int(vol.sum()) # Upscale coarse interior to fine resolution log("Upscaling coarse interior...") interior_up = zoom(vol_c.astype(np.float32), ds_c, order=3) > 0.3 del vol_c gc.collect() # Align and merge (coarse interior may be slightly different shape) for d in range(3): if interior_up.shape[d] + pad_f > grid_f[d]: slc = [slice(None)] * 3 slc[d] = slice(0, grid_f[d] - pad_f) interior_up = interior_up[tuple(slc)] slices = tuple(slice(pad_f, pad_f + interior_up.shape[d]) for d in range(3)) vol[slices] |= interior_up.astype(np.uint8) del interior_up gc.collect() total = int(vol.sum()) log(f"Surface: {surface:,} voxels, with interior: {total:,}") # ── Gaussian smooth + marching cubes ───────────────────────────────── log(f"Gaussian smoothing (sigma={sigma})...") vol = vol.astype(np.float32) gc.collect() gaussian_filter(vol, sigma=sigma, output=vol) # in-place for memory log("Running marching cubes...") verts, faces, normals, _ = measure.marching_cubes(vol, level=0.5) del vol gc.collect() log(f"Marching cubes: {len(verts):,} vertices, {len(faces):,} faces") # Scale back to model coordinates verts = (verts - pad_f) + c_min verts = verts * voxel_size tm = trimesh.Trimesh(vertices=verts, faces=faces, process=True) # ── Keep largest connected component ───────────────────────────────── adj = tm.face_adjacency nf = len(tm.faces) if len(adj) > 0: row = np.concatenate([adj[:, 0], adj[:, 1]]) col = np.concatenate([adj[:, 1], adj[:, 0]]) graph = coo_matrix( (np.ones(len(row), dtype=np.int32), (row, col)), shape=(nf, nf) ) nc, labels = connected_components(graph, directed=False) if nc > 1: from collections import Counter largest = Counter(labels).most_common(1)[0][0] keep = np.where(labels == largest)[0] tm = tm.submesh([keep], append=True) log(f"Kept largest of {nc} components") # ── Mesh smoothing ─────────────────────────────────────────────────── if taubin_iterations > 0: log(f"Taubin mesh smoothing ({taubin_iterations} iterations)...") trimesh.smoothing.filter_taubin(tm, iterations=taubin_iterations) # ── Scale to physical size ─────────────────────────────────────────── if target_height_mm > 0: scale_factor = target_height_mm / tm.extents.max() tm.apply_scale(scale_factor) log(f"Scaled to {target_height_mm:.0f}mm " f"({tm.extents[0]:.1f} x {tm.extents[1]:.1f} x {tm.extents[2]:.1f})") tm.fix_normals() log(f"Watertight: {tm.is_watertight}") log(f"Final: {len(tm.vertices):,} vertices, {len(tm.faces):,} faces") log(f"Completed in {time.time() - t0:.1f}s") return tm # --------------------------------------------------------------------------- # Convenience: run mesh reconstruction in a subprocess (avoids OOM) # --------------------------------------------------------------------------- def voxel_to_mesh_subprocess( mesh_output, output_path: str, target_height_mm: float = 100.0, sigma: float = 1.5, timeout: int = 600, verbose: bool = True, ) -> bool: """ Run voxel_to_mesh in a separate process to avoid OOM. TRELLIS.2 inference retains significant CPU memory even after the model is deleted. Running the mesh reconstruction in a subprocess starts with clean memory, which is necessary for the full-resolution (ds=1) grid on systems with <= 32GB RAM. Args: mesh_output: TRELLIS.2 mesh object. output_path: Where to save the mesh (STL, OBJ, PLY, etc.). target_height_mm: Scale output height in mm. sigma: Gaussian smoothing sigma. timeout: Max seconds for subprocess. verbose: Print progress. Returns: True if the output file was created successfully. """ import torch import subprocess import tempfile # Save coords to temp file if isinstance(mesh_output.coords, torch.Tensor): coords_np = mesh_output.coords.cpu().numpy().copy() else: coords_np = np.array(mesh_output.coords).copy() voxel_size = float(mesh_output.voxel_size) coords_file = tempfile.mktemp(suffix='_coords.npy') np.save(coords_file, coords_np) del coords_np try: result = subprocess.run([ sys.executable, '-c', f""" import numpy as np from blackwell_fix import voxel_coords_to_mesh coords = np.load({coords_file!r}) tm = voxel_coords_to_mesh(coords, {voxel_size}, target_height_mm={target_height_mm}, sigma={sigma}) tm.export({output_path!r}) print(f"Saved: {{output_path}}") """, ], timeout=timeout) return os.path.exists(output_path) finally: try: os.remove(coords_file) except OSError: pass # --------------------------------------------------------------------------- # CLI entry point for standalone testing # --------------------------------------------------------------------------- if __name__ == "__main__": import argparse parser = argparse.ArgumentParser( description="Convert saved TRELLIS.2 voxel coords to watertight mesh" ) parser.add_argument("coords_file", help="Path to .npy file with voxel coordinates") parser.add_argument("voxel_size", type=float, help="Voxel size from mesh.voxel_size") parser.add_argument("-o", "--output", default="output.stl", help="Output mesh path") parser.add_argument("--height", type=float, default=100.0, help="Target height in mm") parser.add_argument("--sigma", type=float, default=1.5, help="Gaussian smoothing sigma") parser.add_argument("--taubin", type=int, default=50, help="Taubin smoothing iterations") args = parser.parse_args() coords = np.load(args.coords_file) tm = voxel_coords_to_mesh( coords, args.voxel_size, target_height_mm=args.height, sigma=args.sigma, taubin_iterations=args.taubin, ) tm.export(args.output) size_mb = os.path.getsize(args.output) / 1e6 print(f"\nSaved {args.output} ({size_mb:.1f} MB)")