From 29ba09b7d5ca57668c0d2d941d4468224b3d0aab Mon Sep 17 00:00:00 2001 From: Noah Peterson Date: Fri, 1 May 2026 13:43:47 -0500 Subject: [PATCH] patches/mps_compat.py: re-enable decode-time cumesh ops on Apple Silicon MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Now that pedronaugusto/mtlmesh#1 ships an Apple7/8 (M1/M2) atomic-fallback for the simplify and atlas kernels, cumesh runs end-to-end on every Apple Silicon GPU family — including the M1 hardware where it previously crashed with: [MtlMesh] Failed to create pipeline for 'propagate_cost_kernel': Unsupported float atomic operation for given target. This previously forced the patcher to unconditionally skip cumesh's three decode-time ops (fill_holes, remove_faces, simplify) on Apple Silicon. Skipping fill_holes specifically had a measurable visual cost: the decoder mesh ships with ~4128 boundary loops worth of small holes, and the bake-time Metal stack only partially closes them after decimation, so the output GLB rendered with see-through cylinders / mesh holes in any glTF viewer. Changes: - patch_mesh_base() no longer inserts an unconditional `return # Skip` at the top of fill_holes / remove_faces / simplify. Instead it: * Replaces `.cuda()` with `.to(self.device)` (the original bug — bare `.cuda()` crashes on MPS regardless of cumesh availability). * Adds an `if cumesh is None: return` guard so CPU-only builds without the Metal stack still skip cleanly. * Leaves the actual cumesh body intact so it runs normally when cumesh is importable. - new patch_pipeline_runtime_fallback() wraps the decode-time `m.fill_holes()` call in pipelines/trellis2_image_to_3d.py in a try/except RuntimeError. Users on unpatched mtlmesh (Apple7/8 without the atomic-fallback fix) get a one-time `warnings.warn` and degrade gracefully to "ship mesh with small holes" rather than crashing the whole pipeline. This is the safety net for users who pip-install before the upstream mtlmesh PR lands. Validation on Apple M1 Pro 16 GB (Apple7) with a patched mtlmesh: - cumesh.fill_holes runs in <1s on a 1,438,207-vertex / 2,988,108-face decoder mesh, closing 4128 boundary loops and adding ~26K faces. - Final GLB face count: 994,484 vs upstream CUDA's 997,498 — within 0.3%. - Visual: see-through artifacts gone; renders match upstream upstream_ref.glb in donmccurdy.com/three.js viewer. Without an mtlmesh patch (legacy install), pipeline still completes; output ships with the original small holes and a one-line warning. Depends on: pedronaugusto/mtlmesh#1 --- patches/mps_compat.py | 76 +++++++++++++++++++++++++++++++++++-------- 1 file changed, 63 insertions(+), 13 deletions(-) diff --git a/patches/mps_compat.py b/patches/mps_compat.py index 732e074..6436bc8 100644 --- a/patches/mps_compat.py +++ b/patches/mps_compat.py @@ -232,11 +232,24 @@ def patch_birefnet(): def patch_mesh_base(): - """Guard cumesh/flex_gemm imports and unconditionally skip in-place mesh - ops. TRELLIS.2 calls fill_holes/remove_faces/simplify during decode on the - full 400K-vertex mesh; the Metal port of cumesh segfaults on inputs that - large, so we skip these decode-time ops entirely. Post-decode mesh - simplification happens later via fast_simplification before texture bake. + """Guard cumesh/flex_gemm imports, replace .cuda() with .to(self.device), + and gate cumesh-using mesh ops on `cumesh is not None`. + + Historical note: this patcher used to unconditionally skip + fill_holes / remove_faces / simplify because Pedro Naugusto's Metal port + of cumesh used `atomic_min`/`atomic_max` on `float`, which is an + Apple9-only feature; on Apple7/8 (M1/M2) the kernel failed at + pipeline-creation time and crashed the whole pipeline. + + With the Apple7/8 atomic-fallback patch in pedronaugusto/mtlmesh#1, + cumesh runs end-to-end on M1/M2, and skipping fill_holes measurably + hurts output quality (decoder mesh ships with ~4128 boundary loops + worth of small holes that the bake-time Metal stack only partially + closes — visible as see-through cylinders in the rendered GLB). + + Soft-failure for users on unpatched mtlmesh is handled at the *caller* + site (`pipelines/trellis2_image_to_3d.py:474`) via a try/except, so + the body of each mesh method here stays straightforward. """ path = os.path.join(TRELLIS_ROOT, "trellis2/representations/mesh/base.py") src = read_file(path) @@ -260,35 +273,35 @@ def patch_mesh_base(): ' raise RuntimeError("flex_gemm requires CUDA")', ) - # Unconditionally return from fill_holes (Metal cumesh segfaults on large meshes) + # Replace `.cuda()` with `.to(self.device)` and add a `cumesh is None` early + # return so methods are no-ops on CPU-only builds rather than crashing. src = src.replace( " def fill_holes(self, max_hole_perimeter=3e-2):\n" " vertices = self.vertices.cuda()\n" " faces = self.faces.cuda()", " def fill_holes(self, max_hole_perimeter=3e-2):\n" - " return # Skip — Metal cumesh segfaults on large decode meshes\n" + " if cumesh is None:\n" + " return # cumesh unavailable (CPU-only build)\n" " vertices = self.vertices.to(self.device)\n" " faces = self.faces.to(self.device)", ) - - # Unconditionally return from remove_faces src = src.replace( " def remove_faces(self, face_mask: torch.Tensor):\n" " vertices = self.vertices.cuda()\n" " faces = self.faces.cuda()", " def remove_faces(self, face_mask: torch.Tensor):\n" - " return\n" + " if cumesh is None:\n" + " return # cumesh unavailable (CPU-only build)\n" " vertices = self.vertices.to(self.device)\n" " faces = self.faces.to(self.device)", ) - - # Unconditionally return from simplify src = src.replace( " def simplify(self, target=1000000, verbose: bool=False, options: dict={}):\n" " vertices = self.vertices.cuda()\n" " faces = self.faces.cuda()", " def simplify(self, target=1000000, verbose: bool=False, options: dict={}):\n" - " return\n" + " if cumesh is None:\n" + " return # cumesh unavailable (CPU-only build)\n" " vertices = self.vertices.to(self.device)\n" " faces = self.faces.to(self.device)", ) @@ -296,6 +309,42 @@ def patch_mesh_base(): write_file(path, src) +def patch_pipeline_runtime_fallback(): + """Wrap decode-time `m.fill_holes()` in a try/except so users on + unpatched mtlmesh (Apple7/8 without the atomic-fallback fix) degrade + gracefully instead of crashing the whole pipeline. + """ + path = os.path.join(TRELLIS_ROOT, "trellis2/pipelines/trellis2_image_to_3d.py") + src = read_file(path) + if "fill_holes_warning_emitted" in src: + print(f" Already patched: {os.path.relpath(path, TRELLIS_ROOT)}") + return + + needle = " m.fill_holes()" + replacement = ( + " try:\n" + " m.fill_holes()\n" + " except RuntimeError as e:\n" + " # Apple7/8 (M1/M2) Metal cumesh without\n" + " # pedronaugusto/mtlmesh#1 raises here. Degrade\n" + " # gracefully (output mesh will have small holes).\n" + " import warnings\n" + " if not getattr(self, '_fill_holes_warning_emitted', False):\n" + " warnings.warn(\n" + ' f"cumesh.fill_holes failed ({e}); skipping. "\n' + ' \"On M1/M2 update mtlmesh to a build with \"\n' + ' \"pedronaugusto/mtlmesh#1 (Apple7/8 atomic fallback).\",\n' + " stacklevel=2,\n" + " )\n" + " self._fill_holes_warning_emitted = True\n" + ) + if needle not in src: + print(f" Skipping (call site moved): {os.path.relpath(path, TRELLIS_ROOT)}") + return + src = src.replace(needle, replacement) + write_file(path, src) + + def patch_fdg_vae(): """Force our pure-Python flexible_dual_grid_to_mesh over any installed o_voxel. The Metal-port o_voxel.convert segfaults on decoder output even @@ -437,6 +486,7 @@ def main(): patch_image_feature_extractor() patch_birefnet() patch_mesh_base() + patch_pipeline_runtime_fallback() patch_fdg_vae() patch_pipeline() patch_pipeline_base()