Fix NaN handling in flash_extract_geometry: preserve NaN for DiffDMC, use -1.0 for MC fallback

patches/triposg/triposg/inference_utils.py

@@ -479,15 +479,19 @@ def flash_extract_geometry(
     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)
+    # Pre-move to CPU float32 while CUDA context is still clean.
+    # Preserve NaN — DiffDMC handles NaN; marching cubes uses a separate nan-filled copy.
+    grid_logits_cpu = grid_logits.float().cpu().numpy()  # preserves NaN
+    device = grid_logits.device
+    del grid_logits
     torch.cuda.empty_cache()
     mesh_v_f = []
     try:
         if DiffDMC is not None:
             print("[TripoSG] Attempting DiffDMC extraction...")
-            gl_gpu = torch.from_numpy(grid_logits_cpu).to(grid_logits.device)
-            dmc = DiffDMC(dtype=torch.float32).to(gl_gpu.device)
+            # Re-upload with NaN intact so DiffDMC ignores masked cells
+            gl_gpu = torch.from_numpy(grid_logits_cpu).to(device)
+            dmc = DiffDMC(dtype=torch.float32).to(device)
             sdf = -gl_gpu / octree_resolution
             sdf = sdf.contiguous()
             vertices, faces = dmc(sdf, deform=None, return_quads=False, normalize=False)
@@ -495,17 +499,21 @@ def flash_extract_geometry(
             faces = faces.detach().cpu().numpy()[:, ::-1]
             vertices = vertices / (2 ** octree_depth) * bbox_size + bbox_min
             mesh_v_f = (vertices.astype(np.float32), np.ascontiguousarray(faces))
-            print(f"[TripoSG] DiffDMC succeeded: {len(vertices)} vertices, {len(faces)} faces")
+            print(f"[TripoSG] DiffDMC: {len(vertices)} vertices, {len(faces)} faces")
+            if len(vertices) == 0:
+                raise RuntimeError("DiffDMC returned empty mesh — falling back to marching cubes")
         else:
             raise RuntimeError("DiffDMC unavailable — using marching cubes")
     except Exception as e:
-        print(f"[TripoSG] DiffDMC failed ({e}), falling back to marching cubes")
+        print(f"[TripoSG] DiffDMC path failed ({e}), falling back to marching cubes")
         torch.cuda.empty_cache()
         try:
-            vertices, faces, normals, _ = measure.marching_cubes(grid_logits_cpu, 0, method="lewiner")
+            # NaN → -1 makes masked/unqueried cells "outside" the isosurface
+            gl_mc = np.nan_to_num(grid_logits_cpu, nan=-1.0)
+            vertices, faces, normals, _ = measure.marching_cubes(gl_mc, 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")
+            print(f"[TripoSG] marching cubes: {len(vertices)} vertices")
         except Exception as e2:
             print(f"[TripoSG] marching cubes also failed: {e2}")
             mesh_v_f = (None, None)