trellis2-blackwell-fix

KO | EN

gitlite — search

#javascript #python #hacktoberfest #react #ai #typescript #llm #go #golang #android #machine-learning #rust #deep-learning #linux

trellis2-blackwell-fix

★ 9 Open GitHub ↗

Workaround for broken mesh extraction in Microsoft TRELLIS.2 on NVIDIA Blackwell GPUs (RTX 5070/5080/5090)

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# trellis2-blackwell-fix

Workaround for broken mesh extraction in [Microsoft TRELLIS.2](https://github.com/microsoft/TRELLIS.2) on NVIDIA Blackwell GPUs (RTX 5070/5070 Ti/5080/5090).

---

## TL;DR

TRELLIS.2's `to_glb()` pipeline calls `cumesh.remeshing.remesh_narrow_band_dc` which silently fails on Blackwell (sm_120) GPUs, producing thousands of disconnected mesh fragments instead of a connected surface. This repository provides a drop-in workaround that bypasses the broken CUDA remeshing entirely and reconstructs a watertight mesh from the raw voxel data using marching cubes. The output is suitable for 3D printing and visualization, though surface detail is reduced compared to native CuMesh output on supported GPUs.

## Affected Hardware

Any GPU with NVIDIA Blackwell architecture (compute capability 12.0):

| GPU | Status |
|---|---|
| **RTX 5070** | Affected (sm_120) |
| **RTX 5070 Ti** | Affected (sm_120) — tested |
| **RTX 5080** | Affected (sm_120) |
| **RTX 5090** | Affected (sm_120) |
| RTX 4090, 4080, etc. | Not affected (sm_89) |
| RTX 3090, 3080, etc. | Not affected (sm_86) |

## Symptoms

If you're running TRELLIS.2 on a Blackwell GPU, you'll see one or more of these:

- **GLB/OBJ output looks like a point cloud** — scattered disconnected vertices with no connected surface
- **Tiny file size** — the exported mesh is suspiciously small (a few MB instead of 50-200MB)
- **Shape is visible but no solid surface** — you can see the rough outline of the object but it's made of thousands of tiny, disconnected triangle patches
- **Slicer can't process the mesh** — PrusaSlicer, Bambu Studio, etc. show errors or an empty build plate
- **`trimesh.is_watertight` returns False** with thousands of connected components

Example diagnostic (run after `pipeline.run()`):

```python
import trimesh
import numpy as np
from scipy.sparse import coo_matrix
from scipy.sparse.csgraph import connected_components

mesh = pipeline.run(image)[0]
v = mesh.vertices.detach().cpu().numpy()
f = mesh.faces.detach().cpu().numpy()
tm = trimesh.Trimesh(vertices=v, faces=f, process=False)

adj = tm.face_adjacency
nf = len(tm.faces)
row = np.concatenate([adj[:, 0], adj[:, 1]])
col = np.concatenate([adj[:, 1], adj[:, 0]])
graph = coo_matrix((np.ones(len(row)), (row, col)), shape=(nf, nf))
nc, labels = connected_components(graph, directed=False)
print(f"Connected components: {nc}")
# Blackwell: ~13,000+ components (broken)
# Working GPU: ~1-5 components (normal)
```

## Root Cause

TRELLIS.2's mesh generation pipeline has two stages:

1. **Dual grid mesh extraction** (`flexible_dual_grid_to_mesh` in `o_voxel`) — converts the sparse voxel latent representation into a raw mesh. This raw mesh is intentionally sparse: it consists of thousands of small, disconnected triangle patches covering the surface.

2. **CuMesh remeshing** (`remesh_narrow_band_dc`) — stitches the sparse patches into a single connected mesh with proper topology, then decimates to a target face count.

**Stage 2 silently fails on Blackwell.** The CuMesh remeshing step does not crash — it completes without errors — but produces garbage output. The fragments that should be connected remain disconnected. The resulting mesh has the right rough shape but no coherent surface.

This is **not** a CUDA hashmap bug in the dual grid extraction. We verified this by reimplementing `flexible_dual_grid_to_mesh` entirely in CPU Python/NumPy — the raw mesh is equally fragmented. The fragmentation is inherent to stage 1; it's stage 2 (CuMesh remeshing) that's supposed to fix it and doesn't on sm_120.

The underlying cause is likely that CuMesh's CUDA kernels produce incorrect results on the Blackwell instruction set. The compute capability patch (reporting sm_90 instead of sm_120) enables the kernels to run but doesn't fix the numerical results. This is a problem in the CuMesh library, not in TRELLIS.2 itself.

## Dead Ends (Save Yourself Time)

We tried these approaches before arriving at the current workaround:

### CPU hashmap fallback for `flexible_dual_grid_to_mesh`
**Result:** Same ~14K disconnected fragments. The CPU reimplementation faithfully reproduces the fragmented output, proving the issue is not a CUDA hashmap bug. The raw dual grid mesh is *supposed* to be fragmented — it's the downstream remeshing that's broken.

### Poisson surface reconstruction (Open3D)
**Result:** Crumbly, noisy surface. The raw vertices from the fragmented mesh have inconsistent normals across disconnected patches. Estimated normals (`estimate_normals()`) are also unreliable because the point cloud has no consistent orientation. The Poisson reconstruction produces a blobby, unusable mesh.

### Direct export of `mesh.vertices` / `mesh.faces`
**Result:** The raw mesh has ~760K vertices but only ~91K faces, with ~13K disconnected components. The largest component has only ~24 faces. This is the raw dual grid output before remeshing — it's not meant to be used directly.

## Working Workaround

The workaround bypasses both the broken mesh extraction and CuMesh remeshing entirely. Instead, it uses the **voxel coordinates** (`mesh.coords`) which are reliably produced on all GPUs:

1. **Coarse interior fill** (downsampled 4x) — aggressive morphological dilation seals gaps in the sparse voxel shell, `binary_fill_holes` determines the solid interior, erosion restores the surface position
2. **Full-resolution surface** (ds=1) — original voxel positions combined with upscaled coarse interior for a solid volume
3. **Gaussian smooth + marching cubes** — smooths the binary volume and extracts an isosurface
4. **Largest component extraction** — removes small disconnected artifacts
5. **Taubin mesh smoothing** — reduces voxel staircase artifacts

The result is a watertight mesh suitable for 3D printing. Surface detail is limited to the voxel grid resolution (~0.1mm at typical print scales), which is below FDM layer height but visible in renders.

## Installation

```bash
# Clone this repo
git clone https://github.com/ThatButters/trellis2-blackwell-fix
cd trellis2-blackwell-fix

# Install dependencies (assumes TRELLIS.2 is already set up)
pip install numpy scipy scikit-image trimesh
```

## Usage

### As a drop-in module

```python
import blackwell_fix

# Apply patches BEFORE importing trellis2
blackwell_fix.patch_all()

from trellis2.pipelines import Trellis2ImageTo3DPipeline
from PIL import Image
import torch

# Standard TRELLIS.2 inference
pipeline = Trellis2ImageTo3DPipeline.from_pretrained("microsoft/TRELLIS.2-4B")
pipeline.cuda()
image = Image.open("photo.jpg")
torch.manual_seed(42)
mesh = pipeline.run(image)[0]

# Replace broken to_glb() with voxel workaround
result = blackwell_fix.voxel_to_mesh(mesh, target_height_mm=100.0)
result.export("output.stl")
```

### With subprocess (recommended for <= 32GB RAM)

TRELLIS.2 inference retains significant CPU memory even after the model is deleted. On systems with limited RAM, run the mesh reconstruction in a subprocess:

```python
# After inference, free the model first
coords_np = mesh.coords.cpu().numpy().copy()
voxel_size = float(mesh.voxel_size)
del pipeline, mesh
torch.cuda.empty_cache()

# Run in subprocess with clean memory
import numpy as np
np.save("coords.npy", coords_np)
# Then in a new process:
# python blackwell_fix.py coords.npy 0.000977 -o output.stl
```

### CLI

```bash
# Reconstruct from saved voxel coordinates
python blackwell_fix.py coords.npy 0.000977 -o output.stl --height 100 --sigma 1.5
```

### Parameters

| Parameter | Default | Description |
|---|---|---|
| `target_height_mm` | 100.0 | Scale output so tallest dimension = this (mm). Set to 0 to skip. |
| `sigma` | 1.5 | Gaussian smoothing sigma. Higher = smoother but less detail. |
| `coarse_downsample` | 4 | Downsample factor for interior fill phase. |
| `taubin_iterations` | 50 | Mesh smoothing passes. More = smoother surface. |

## In Progress

- **Multi-resolution hybrid approach**: Using full-resolution surface detail with coarse interior fill. Currently limited by memory (ds=1 grid is ~380M cells / 1.5GB as float32).
- **Intercepting `dual_vertices`**: TRELLIS.2 internally computes float-precision vertex offsets (`dual_vertices`) that contain sub-voxel surface detail. If these can be correctly extracted and used for mesh construction, the quality improvement would be dramatic. See [CONTRIBUTING.md](CONTRIBUTING.md) for details.

## Compatibility

| Component | Tested Version |
|---|---|
| GPU | RTX 5070 Ti (sm_120) |
| PyTorch | 2.7.0+cu128 |
| CUDA | 12.8 |
| OS | Ubuntu 22.04 (WSL2) |
| TRELLIS.2 | 4B model (`microsoft/TRELLIS.2-4B`) |
| Python | 3.10 |
| numpy | 1.26.x (must be < 2.0 for spconv compatibility) |
| scipy | 1.12+ |
| scikit-image | 0.22+ |
| trimesh | 4.0+ |

## Broader Impact

This issue likely affects **any project using CuMesh or O-Voxel remeshing on Blackwell GPUs**, including but not limited to:

- **[ComfyUI-TRELLIS2](https://github.com/JMaiGC/ComfyUI-TRELLIS2)** — popular ComfyUI integration for TRELLIS.2
- **Instant-NGP / Nerfstudio** — NeRF pipelines that use marching cubes + CuMesh for mesh extraction
- **Sparse convolution networks** (spconv, MinkowskiEngine) — the CC detection patches in this repo are needed for these libraries too
- **3D object detection** (OpenPCDet, CenterPoint) — spconv-based models need the same CC patches
- **Any pipeline using `o_voxel.postprocess.to_glb()`** with `remesh=True`

If you maintain one of these projects and are seeing issues on RTX 5070/5080/5090, the patches in `blackwell_fix.patch_all()` may help with the CUDA/spconv/Triton compatibility issues, though the CuMesh remeshing problem will require a separate fix specific to your mesh extraction pipeline.

## Related Issues

- [microsoft/TRELLIS#162](https://github.com/microsoft/TRELLIS/issues/162) — TRELLIS (v1) Blackwell compatibility
- [traveller59/spconv#754](https://github.com/traveller59/spconv/issues/754) — spconv SIGFPE on Blackwell
- [openai/triton#5765](https://github.com/openai/triton/issues/5765) — Triton compilation failure on sm_120

## License

MIT — see [LICENSE](LICENSE).

trellis2-blackwell-fix

Explore Similar Repositories

Was this content helpful?

Recent Feedback

Download README