add zerogpu

inference/inference.py

@@ -2,6 +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
 
 DEVICE = "cpu"
 if torch.cuda.is_available():
@@ -75,14 +76,16 @@ def get_heatmap_rgb(model, data, N_CHUNKS=5): # evaluate loader can only have ba
         scores = all_logits.squeeze().cpu()
         heatmap_rgb = torch.tensor(plt.cm.jet(scores.numpy())[:,:3]).squeeze()
         return heatmap_rgb
-
+    
+@spaces.GPU(duration=90)
 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))
     data_dict["label_embeds"] = encode_text(queries)
     seg_rgb = get_segmentation_rgb(model, data_dict)
     return seg_rgb
-
+    
+@spaces.GPU(duration=90)
 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))