Adapt to ZeroGPU?

app.py

@@ -43,7 +43,7 @@ def _safe_ext(path: str) -> str:
             return ext
     return os.path.splitext(path)[1].lower()
 
-@spaces.GPU
+
 def convert_and_show(mesh_file):
     os.makedirs("tmp/glbs", exist_ok=True)
     if mesh_file is None:
@@ -54,7 +54,6 @@ def convert_and_show(mesh_file):
     mesh.export(f_name)
     return f_name
 
-@spaces.GPU
 def convert_and_show_twice(mesh_file_1, mesh_file_2):
     return convert_and_show(mesh_file_1), convert_and_show(mesh_file_2)
 
@@ -176,6 +175,8 @@ def build_outputs(surf_a: Surface, surf_b: Surface, cmap_a: np.ndarray, p2p: np.
 @spaces.GPU
 def init_clicked(mesh1_path, mesh2_path,
                  lambda_val, zoomout_val, time_val, nloop_val, sds_val, proper_val):
+
+    matcher._init()
     print("inside init")
     cfg.deepfeat_conf.fmap.lambda_ = lambda_val
     cfg.sds_conf.zoomout = zoomout_val
@@ -188,7 +189,6 @@ def init_clicked(mesh1_path, mesh2_path,
         raise gr.Error("Please upload both meshes.")
     global datadicts
     datadicts = Datadicts(mesh1_path, mesh2_path)
-    matcher._init()
     shape_dict, target_dict = convert_dict(datadicts.shape_dict, 'cuda'), convert_dict(datadicts.target_dict, 'cuda')
     fmap_model_cuda = matcher.fmap_model.cuda()
     diff_model_cuda = matcher.diffusion_model

zero_shot.py

@@ -27,6 +27,7 @@ from utils.pickle_stuff import safe_load_with_fallback
 from utils.geometry import compute_operators, load_operators
 from utils.surfaces import Surface
 import sys
+import spaces
 
 try:
     import google.colab
@@ -120,6 +121,7 @@ class Matcher(object):
         self.fmap_cfg = self.cfg.deepfeat_conf.fmap
         self.dataloaders = dict()
 
+    @spaces.GPU
     def _init(self):
         cfg = self.cfg
         self.fmap_model = DFMNet(self.cfg["deepfeat_conf"]["fmap"]).cuda()
@@ -136,6 +138,7 @@ class Matcher(object):
         self.eye = torch.eye(self.n_fmap).float().cuda()
         self.eye.requires_grad = False
 
+    @spaces.GPU
     def fmap(self, shape_dict, target_dict):
         diff_model_cuda = self.diffusion_model
         diff_model_cuda.net.cuda()
@@ -152,6 +155,7 @@ class Matcher(object):
             mask_12, mask_21 = None, None
         return C12_pred, C12_obj, C21_pred, C21_obj, feat1, feat2, evecs_trans1, evecs_trans2, mask_12, mask_21
 
+    @spaces.GPU
     def zo_shot(self, shape_dict, target_dict):
         self._init()
         evecs1, evecs2 = shape_dict["evecs"], target_dict["evecs"]
@@ -161,6 +165,7 @@ class Matcher(object):
         indKNN_new, _ = extract_p2p_torch_fmap(new_FM, evecs1, evecs2)
         return new_FM, indKNN_new
 
+    @spaces.GPU
     def optimize(self, shape_dict, target_dict, target_normals):
         self._init()
         diff_model_cuda = self.diffusion_model