fix open3d zerogpu

app.py

@@ -1,4 +1,6 @@
 import gradio as gr
+import open3d_zerogpu_fix
+import spaces
 import re
 from utils import read_pcd, render_point_cloud, render_pcd_file, set_seed, install_cuda_toolkit
 from inference.utils import get_legend
@@ -43,6 +45,10 @@ source_dict = {
     "chair": "wild"
 }
 
+@spaces.GPU(duration=120)
+def run_predict(*args):
+    yield from predict(*args)
+
 def predict(pcd_path, inference_mode, part_queries):
     set_seed()
     xyz, rgb, normal = read_pcd(pcd_path)
@@ -52,14 +58,14 @@ def predict(pcd_path, inference_mode, part_queries):
             raise gr.Error("For segmentation mode, please provide 2 or more parts", duration=5)
         seg_rgb = segment_obj(xyz, rgb, normal, parts).cpu().numpy()
         legend = get_legend(parts)
-        return render_point_cloud(xyz, seg_rgb, legend=legend)
+        yield render_point_cloud(xyz, seg_rgb, legend=legend)
     elif inference_mode == "Localization":
         if "," in part_queries or ";" in part_queries or "." in part_queries:
             raise gr.Error("For localization mode, please provide only one part", duration=5)
         heatmap_rgb = get_heatmap(xyz, rgb, normal, part_queries).cpu().numpy()
-        return render_point_cloud(xyz, heatmap_rgb)
+        yield render_point_cloud(xyz, heatmap_rgb)
     else:
-        return None
+        yield None
 
 def on_select(evt: gr.SelectData):
     obj_name = evt.value['image']['orig_name'][:-4]
@@ -150,7 +156,7 @@ with gr.Blocks(theme=gr.themes.Default(text_size="lg", radius_size="none")) as d
         outputs=[input_point_cloud],
     )
     run_button.click(
-        fn=predict,
+        fn=run_predict,
         inputs=[file_upload, inference_mode, part_queries],
         outputs=[output_point_cloud],
     )

inference/inference.py

@@ -2,7 +2,7 @@ import torch
 import numpy as np
 import matplotlib.pyplot as plt
 from inference.utils import get_seg_color, load_model, preprocess_pcd, encode_text
-import spaces
+#import spaces
 
 DEVICE = "cuda:0"
 #if torch.cuda.is_available():
@@ -77,7 +77,7 @@ def get_heatmap_rgb(model, data, N_CHUNKS=5): # evaluate loader can only have ba
         heatmap_rgb = torch.tensor(plt.cm.jet(scores.numpy())[:,:3]).squeeze()
         return heatmap_rgb
     
-@spaces.GPU
+#@spaces.GPU
 def segment_obj(xyz, rgb, normal, queries):
     model = load_model()
     data_dict = preprocess_pcd(torch.tensor(xyz).float().to(DEVICE), torch.tensor(rgb).float().to(DEVICE), torch.tensor(normal).float().to(DEVICE))
@@ -85,7 +85,7 @@ def segment_obj(xyz, rgb, normal, queries):
     seg_rgb = get_segmentation_rgb(model, data_dict)
     return seg_rgb
     
-@spaces.GPU
+#@spaces.GPU
 def get_heatmap(xyz, rgb, normal, query):
     model = load_model()
     data_dict = preprocess_pcd(torch.tensor(xyz).float().to(DEVICE), torch.tensor(rgb).float().to(DEVICE), torch.tensor(normal).float().to(DEVICE))

open3d_zerogpu_fix.py

@@ -0,0 +1,7 @@
+import fileinput
+import site
+from pathlib import Path
+
+with fileinput.FileInput(f'{site.getsitepackages()[0]}/open3d/__init__.py', inplace=True) as file:
+    for line in file:
+        print(line.replace('_pybind_cuda.open3d_core_cuda_device_count()', '1'), end='')