Adapt to ZeroGPU?

app.py

@@ -193,7 +193,7 @@ def init_clicked(mesh1_path, mesh2_path,
     diff_model_cuda.net.cuda()
     C12_pred_init, C21_pred_init, feat1, feat2, evecs_trans1, evecs_trans2 = fmap_model_cuda({"shape1": shape_dict, "shape2": target_dict}, diff_model=diff_model_cuda, scale=matcher.fmap_cfg.diffusion.time)
     C12_pred, C12_obj, mask_12 = C12_pred_init
-    p2p_init, _ = extract_p2p_torch_fmap(C12_obj, datadicts.shape_dict["evecs"], datadicts.target_dict["evecs"])
+    p2p_init, _ = extract_p2p_torch_fmap(C12_obj, shape_dict["evecs"], target_dict["evecs"])
     return build_outputs(datadicts.shape_surf, datadicts.target_surf, datadicts.cmap1, p2p_init, tag="init")
 
 @spaces.GPU
@@ -218,14 +218,15 @@ def run_clicked(mesh1_path, mesh2_path, yaml_path, lambda_val, zoomout_val, time
         if not (datadicts.shape_path == mesh1_path and datadicts.target_path == mesh2_path):
             datadicts = Datadicts(mesh1_path, mesh2_path)
 
+    shape_dict, target_dict = convert_dict(datadicts.shape_dict, 'cuda'), convert_dict(datadicts.target_dict, 'cuda')
     target_normals = torch.from_numpy(datadicts.target_surf.surfel/np.linalg.norm(datadicts.target_surf.surfel, axis=-1, keepdims=True)).float().to("cuda")
 
-    C12_new, p2p, p2p_init, _, loss_save = matcher.optimize(datadicts.shape_dict, datadicts.target_dict, target_normals)
-    evecs1, evecs2 = datadicts.shape_dict["evecs"], datadicts.target_dict["evecs"]
-    evecs_2trans = evecs2.t() @ torch.diag(datadicts.target_dict["mass"])
+    C12_new, p2p, p2p_init, _, loss_save = matcher.optimize(shape_dict, target_dict, target_normals)
+    evecs1, evecs2 = shape_dict["evecs"], target_dict["evecs"]
+    evecs_2trans = evecs2.t() @ torch.diag(target_dict["mass"])
     with torch.no_grad():
         C12_end_zo = torch_zoomout(evecs1, evecs2, evecs_2trans, C12_new.squeeze(), matcher.cfg.sds_conf.zoomout)
-    p2p_zo, _ = extract_p2p_torch_fmap(C12_end_zo, datadicts.shape_dict["evecs"], datadicts.target_dict["evecs"])
+    p2p_zo, _ = extract_p2p_torch_fmap(C12_end_zo, shape_dict["evecs"], target_dict["evecs"])
     return build_outputs(datadicts.shape_surf, datadicts.target_surf, datadicts.cmap1, p2p_zo, tag="run")
 
 
@@ -319,7 +320,7 @@ if __name__ == "__main__":
     print("Making matcher")
     matcher = zero_shot.Matcher(cfg)
     print("Matcher ready")
-    shutil.rmtree("tmp")
+    #shutil.rmtree("tmp")
     os.makedirs("tmp", exist_ok=True)
     datadicts = None
     demo.launch(share=args.share)