Remove pykeops for space ZeroGPU demo.

app.py

@@ -30,6 +30,7 @@ from utils.meshplot import visu_pts
 from utils.torch_fmap import extract_p2p_torch_fmap, torch_zoomout
 import torch
 import argparse
+from utils.utils_func import convert_dict
 # -----------------------------
 # Utils
 # -----------------------------
@@ -173,8 +174,11 @@ def init_clicked(mesh1_path, mesh2_path,
     matcher._init()
     global datadicts
     datadicts = Datadicts(mesh1_path, mesh2_path)
-
-    C12_pred_init, C21_pred_init, feat1, feat2, evecs_trans1, evecs_trans2 = matcher.fmap_model({"shape1": datadicts.shape_dict, "shape2": datadicts.target_dict}, diff_model=matcher.diffusion_model, scale=matcher.fmap_cfg.diffusion.time)
+    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
+    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"])
     return build_outputs(datadicts.shape_surf, datadicts.target_surf, datadicts.cmap1, p2p_init, tag="init")

zero_shot.py

@@ -100,11 +100,9 @@ class Matcher(object):
 
     def __init__(self, cfg):
         self.cfg = cfg
-        self.device = torch.device(f'cuda:{cfg["gpu"]}' if torch.cuda.is_available() else 'cpu')
-        print(f"Using device: {self.device}")
         self.diffusion_model = None
         if self.cfg.get("sds", False):
-            self.diffusion_model = DiffModel(cfg["sds_conf"], self.device)
+            self.diffusion_model = DiffModel(cfg["sds_conf"], "cpu")
         self.n_fmap = self.cfg["deepfeat_conf"]["fmap"]["n_fmap"]
         self.n_loop = 0
         if self.cfg.get("optimize", False):
@@ -124,10 +122,10 @@ class Matcher(object):
 
     def _init(self):
         cfg = self.cfg
-        self.fmap_model = DFMNet(self.cfg["deepfeat_conf"]["fmap"]).to(self.device)
+        self.fmap_model = DFMNet(self.cfg["deepfeat_conf"]["fmap"]).cuda()
         if self.snk:
-            self.encoder = Encoder().to(self.device)
-            self.decoder = PrismDecoder(dim_in=515).to(self.device)
+            self.encoder = Encoder().cuda()
+            self.decoder = PrismDecoder(dim_in=515).cuda()
             self.loss_prism = PrismRegularizationLoss(primo_h=0.02)
             self.soft_p2p = True
             params_to_opt = list(self.fmap_model.parameters()) + list(self.encoder.parameters()) + list(
@@ -135,13 +133,15 @@ class Matcher(object):
         else:
             params_to_opt = self.fmap_model.parameters()
         self.optim = torch.optim.Adam(params_to_opt, lr=0.001, betas=(0.9, 0.99))
-        self.eye = torch.eye(self.n_fmap).float().to(self.device)
+        self.eye = torch.eye(self.n_fmap).float().cuda()
         self.eye.requires_grad = False
 
     def fmap(self, shape_dict, target_dict):
+        diff_model_cuda = self.diffusion_model
+        diff_model_cuda.net.cuda()
         if self.fmap_cfg.get("use_diff", False):
             C12_pred, C21_pred, feat1, feat2, evecs_trans1, evecs_trans2 = self.fmap_model(
-                {"shape1": shape_dict, "shape2": target_dict}, diff_model=self.diffusion_model,
+                {"shape1": shape_dict, "shape2": target_dict}, diff_model=diff_model_cuda,
                 scale=self.fmap_cfg.diffusion.time)
             C12_pred, C12_obj, mask_12 = C12_pred
             C21_pred, C21_obj, mask_21 = C21_pred
@@ -163,6 +163,8 @@ class Matcher(object):
 
     def optimize(self, shape_dict, target_dict, target_normals):
         self._init()
+        diff_model_cuda = self.diffusion_model
+        diff_model_cuda.net.cuda()
         evecs1, evecs2 = shape_dict["evecs"], target_dict["evecs"]
         C12_pred_init, _, _, _, _, _, evecs_trans1, evecs_trans2, _, _ = self.fmap(shape_dict, target_dict)
         evecs_2trans = evecs2.t() @ torch.diag(target_dict["mass"])
@@ -203,8 +205,7 @@ class Matcher(object):
             l_lap += ((C21_new @ torch.diag(target_dict["evals"][:self.n_fmap]) - torch.diag(
                 shape_dict["evals"][:self.n_fmap]) @ C21_new) ** 2).mean()
 
-            l_cycle, l_prism, l_mse = torch.as_tensor(0.).float().to(self.device), torch.as_tensor(0.).float().to(
-                self.device), torch.as_tensor(0.).float().to(self.device)
+            l_cycle, l_prism, l_mse = torch.as_tensor(0.).float().cuda(), torch.as_tensor(0.).float().cuda(), torch.as_tensor(0.).float().cuda()
             if self.snk:
                 # Latent vector
                 latents = self.encoder(shape_dict)
@@ -217,7 +218,7 @@ class Matcher(object):
                 l_mse = ((soft_p2p_21 @ shape_dict["vertices"] - snk_rec) ** 2).sum(dim=-1).mean()
                 l_cycle = ((soft_p2p_12 @ (soft_p2p_21 @ shape_dict["vertices"]) - shape_dict["vertices"]) ** 2).sum(
                     dim=-1).mean()
-            l_sds, l_proper = torch.as_tensor(0.).float().to(self.device), torch.as_tensor(0.).float().to(self.device)
+            l_sds, l_proper = torch.as_tensor(0.).float().cuda(), torch.as_tensor(0.).float().cuda()
             if self.fmap_cfg.get("use_diff", False):
                 if self.fmap_cfg.diffusion.get("abs", False):
                     C12_in, C21_in = torch.abs(C12_pred).squeeze(), torch.abs(C21_pred).squeeze()
@@ -225,7 +226,7 @@ class Matcher(object):
                     C12_in, C21_in = C12_pred.squeeze(), C21_pred.squeeze()
                 grad_12, _ = guidance_grad(C12_in, self.diffusion_model.net, grad_scale=1,
                                            batch_size=self.fmap_cfg.diffusion.batch_sds,
-                                           scale_noise=self.fmap_cfg.diffusion.time, device=self.device)
+                                           scale_noise=self.fmap_cfg.diffusion.time, device="cuda")
                 with torch.no_grad():
                     denoised_12 = C12_pred - self.optim.param_groups[0]['lr'] * grad_12
                 targets_12 = torch_zoomout(evecs1, evecs2, evecs_2trans, C12_obj.squeeze(), self.cfg.sds_conf.zoomout)
@@ -233,9 +234,9 @@ class Matcher(object):
                 l_proper_12 = ((C12_pred.squeeze()[:self.n_fmap, :self.n_fmap] - targets_12.squeeze()[:self.n_fmap,
                                                                                  :self.n_fmap]) ** 2).mean()
 
-                grad_21, _ = guidance_grad(C21_in, self.diffusion_model.net, grad_scale=1,
+                grad_21, _ = guidance_grad(C21_in, diff_model_cuda.net, grad_scale=1,
                                            batch_size=self.fmap_cfg.diffusion.batch_sds,
-                                           scale_noise=self.fmap_cfg.diffusion.time, device=self.device)
+                                           scale_noise=self.fmap_cfg.diffusion.time, device="cuda")
                 # denoised_21 = C21_pred - self.optim.param_groups[0]['lr'] * grad_21
                 with torch.no_grad():
                     denoised_21 = C21_pred - self.optim.param_groups[0]['lr'] * grad_21
@@ -251,7 +252,7 @@ class Matcher(object):
                 l_sds += ((torch.abs(C21_pred).squeeze()[:self.n_fmap, :self.n_fmap] - denoised_21.squeeze()[
                                                                                        :self.n_fmap,
                                                                                        :self.n_fmap]) ** 2).mean()
-            loss = torch.as_tensor(0.).float().to(self.device)
+            loss = torch.as_tensor(0.).float().to("cuda")
             if self.cfg.loss.get("ortho", 0) > 0:
                 loss += self.cfg.loss.get("ortho", 0) * l_ortho
             if self.cfg.loss.get("bij", 0) > 0: