Pre-convert grid_logits to CPU before DiffDMC attempt (prevents CUDA ctx corruption in fallback)

patches/triposg/triposg/inference_utils.py

@@ -478,13 +478,18 @@ def flash_extract_geometry(
     torch.cuda.empty_cache()
     mesh_v_f = []
     grid_logits = grid_logits[0]
+    print("final grids shape = ", grid_logits.shape)
+    # Pre-convert to CPU float32 before any CUDA-compiled ops (DiffDMC may corrupt CUDA ctx)
+    grid_logits_cpu = np.nan_to_num(grid_logits.float().cpu().numpy(), nan=0.0)
+    torch.cuda.empty_cache()
+    mesh_v_f = []
     try:
-        print("final grids shape = ", grid_logits.shape)
         if DiffDMC is not None:
             print("[TripoSG] Attempting DiffDMC extraction...")
-            dmc = DiffDMC(dtype=torch.float32).to(grid_logits.device)
-            sdf = -grid_logits / octree_resolution
-            sdf = sdf.to(torch.float32).contiguous()
+            gl_gpu = torch.from_numpy(grid_logits_cpu).to(grid_logits.device)
+            dmc = DiffDMC(dtype=torch.float32).to(gl_gpu.device)
+            sdf = -gl_gpu / octree_resolution
+            sdf = sdf.contiguous()
             vertices, faces = dmc(sdf, deform=None, return_quads=False, normalize=False)
             vertices = vertices.detach().cpu().numpy()
             faces = faces.detach().cpu().numpy()[:, ::-1]
@@ -492,28 +497,17 @@ def flash_extract_geometry(
             mesh_v_f = (vertices.astype(np.float32), np.ascontiguousarray(faces))
             print(f"[TripoSG] DiffDMC succeeded: {len(vertices)} vertices, {len(faces)} faces")
         else:
-            # DiffDMC unavailable — fall back to marching cubes
-            print("[TripoSG] DiffDMC unavailable, falling back to marching cubes")
-            gl = grid_logits.float().cpu().numpy()
-            gl_clean = np.nan_to_num(gl, nan=0.0)
-            vertices, faces, normals, _ = measure.marching_cubes(gl_clean, 0, method="lewiner")
-            vertices = vertices / (2 ** octree_depth) * bbox_size + bbox_min
-            mesh_v_f = (vertices.astype(np.float32), np.ascontiguousarray(faces))
-            print(f"[TripoSG] marching cubes succeeded: {len(vertices)} vertices")
+            raise RuntimeError("DiffDMC unavailable — using marching cubes")
     except Exception as e:
-        import traceback as _tb
-        print(f"[TripoSG] mesh extraction failed: {e}")
-        print(_tb.format_exc())
+        print(f"[TripoSG] DiffDMC failed ({e}), falling back to marching cubes")
         torch.cuda.empty_cache()
-        # Last-resort: marching cubes on nan-cleaned grid
         try:
-            gl_np = np.nan_to_num(grid_logits.float().cpu().numpy(), nan=0.0)
-            vertices, faces, normals, _ = measure.marching_cubes(gl_np, 0, method="lewiner")
-            vertices = vertices / max(gl_np.shape) * bbox_size + bbox_min
+            vertices, faces, normals, _ = measure.marching_cubes(grid_logits_cpu, 0, method="lewiner")
+            vertices = vertices / (2 ** octree_depth) * bbox_size + bbox_min
             mesh_v_f = (vertices.astype(np.float32), np.ascontiguousarray(faces))
-            print(f"[TripoSG] last-resort MC succeeded: {len(vertices)} vertices")
+            print(f"[TripoSG] marching cubes succeeded: {len(vertices)} vertices")
         except Exception as e2:
-            print(f"[TripoSG] last-resort MC also failed: {e2}")
+            print(f"[TripoSG] marching cubes also failed: {e2}")
             mesh_v_f = (None, None)
 
     return [mesh_v_f]