Fix queries dtype cast in inference_utils.py; remove autocast (caused NaN/no isosurface)

app.py

@@ -216,7 +216,7 @@ def load_triposg():
     _iu_path = triposg_src / "triposg" / "inference_utils.py"
     if _iu_path.exists():
         _iu_text = _iu_path.read_text()
-        if "DiffDMC = None" not in _iu_text and "from diso import DiffDMC" in _iu_text:
+        if "queries.to(dtype=batch_latents.dtype)" not in _iu_text:
             _iu_text = _iu_text.replace(
                 "from diso import DiffDMC",
                 "try:\n    from diso import DiffDMC\n"
@@ -237,8 +237,15 @@ def load_triposg():
                     "        return flash_extract_geometry(*args, **kwargs)\n"
                     "    return _hierarchical_extract_geometry_impl(*args, **kwargs)\n"
                 )
+            # Also cast queries to match batch_latents dtype before vae.decode.
+            # TripoSGPipeline loads as float16 but flash_extract_geometry creates
+            # query grids as float32, causing a dtype mismatch in F.linear.
+            _iu_text = _iu_text.replace(
+                "logits = vae.decode(batch_latents, queries).sample",
+                "logits = vae.decode(batch_latents, queries.to(dtype=batch_latents.dtype)).sample",
+            )
             _iu_path.write_text(_iu_text)
-            print("[load_triposg] Patched inference_utils.py: diso optional")
+            print("[load_triposg] Patched inference_utils.py: diso optional + queries dtype cast")
 
     weights_path = snapshot_download("VAST-AI/TripoSG")
 
@@ -334,16 +341,15 @@ def generate_shape(input_image, remove_background, num_steps, guidance_scale,
 
         progress(0.5, desc="Generating shape (SDF diffusion)...")
         from scripts.inference_triposg import run_triposg
-        with torch.autocast(device_type="cuda", dtype=torch.float16):
-            mesh = run_triposg(
-                pipe=pipe,
-                image_input=img_path,
-                rmbg_net=rmbg_net if remove_background else None,
-                seed=int(seed),
-                num_inference_steps=int(num_steps),
-                guidance_scale=float(guidance_scale),
-                faces=int(face_count) if int(face_count) > 0 else -1,
-            )
+        mesh = run_triposg(
+            pipe=pipe,
+            image_input=img_path,
+            rmbg_net=rmbg_net if remove_background else None,
+            seed=int(seed),
+            num_inference_steps=int(num_steps),
+            guidance_scale=float(guidance_scale),
+            faces=int(face_count) if int(face_count) > 0 else -1,
+        )
 
         out_path = "/tmp/triposg_shape.glb"
         mesh.export(out_path)