add cpu-only mode

cotracker/evaluation/core/evaluator.py

@@ -185,7 +185,11 @@ class Evaluator:
                 if not all(gotit):
                     print("batch is None")
                     continue
-            dataclass_to_cuda_(sample)
+            if torch.cuda.is_available():
+                dataclass_to_cuda_(sample)
+                device = torch.device("cuda")
+            else:
+                device = torch.device("cpu")
 
             if (
                 not train_mode
@@ -205,7 +209,7 @@ class Evaluator:
                         queries[:, :, 1],
                     ],
                     dim=2,
-                )
+                ).to(device)
             else:
                 queries = torch.cat(
                     [
@@ -213,7 +217,7 @@ class Evaluator:
                         sample.trajectory[:, 0],
                     ],
                     dim=2,
-                )
+                ).to(device)
 
             pred_tracks = model(sample.video, queries)
             if "strided" in dataset_name:

cotracker/evaluation/evaluate.py

@@ -102,6 +102,8 @@ def run_eval(cfg: DefaultConfig):
         single_point=cfg.single_point,
         n_iters=cfg.n_iters,
     )
+    if torch.cuda.is_available():
+        predictor.model = predictor.model.cuda()
 
     # Setting the random seeds
     torch.manual_seed(cfg.seed)

cotracker/models/core/cotracker/cotracker.py

@@ -25,14 +25,14 @@ from cotracker.models.core.embeddings import (
 torch.manual_seed(0)
 
 
-def get_points_on_a_grid(grid_size, interp_shape, grid_center=(0, 0)):
+def get_points_on_a_grid(grid_size, interp_shape, grid_center=(0, 0), device="cuda"):
     if grid_size == 1:
-        return torch.tensor([interp_shape[1] / 2, interp_shape[0] / 2])[
+        return torch.tensor([interp_shape[1] / 2, interp_shape[0] / 2], device=device)[
             None, None
-        ].cuda()
+        ]
 
     grid_y, grid_x = meshgrid2d(
-        1, grid_size, grid_size, stack=False, norm=False, device="cuda"
+        1, grid_size, grid_size, stack=False, norm=False, device=device
     )
     step = interp_shape[1] // 64
     if grid_center[0] != 0 or grid_center[1] != 0:
@@ -47,7 +47,7 @@ def get_points_on_a_grid(grid_size, interp_shape, grid_center=(0, 0)):
 
     grid_y = grid_y + grid_center[0]
     grid_x = grid_x + grid_center[1]
-    xy = torch.stack([grid_x, grid_y], dim=-1).cuda()
+    xy = torch.stack([grid_x, grid_y], dim=-1).to(device)
     return xy
 
 

cotracker/models/evaluation_predictor.py

@@ -29,11 +29,10 @@ class EvaluationPredictor(torch.nn.Module):
         self.n_iters = n_iters
 
         self.model = cotracker_model
-        self.model.to("cuda")
         self.model.eval()
 
     def forward(self, video, queries):
-        queries = queries.clone().cuda()
+        queries = queries.clone()
         B, T, C, H, W = video.shape
         B, N, D = queries.shape
 
@@ -42,14 +41,16 @@ class EvaluationPredictor(torch.nn.Module):
 
         rgbs = video.reshape(B * T, C, H, W)
         rgbs = F.interpolate(rgbs, tuple(self.interp_shape), mode="bilinear")
-        rgbs = rgbs.reshape(B, T, 3, self.interp_shape[0], self.interp_shape[1]).cuda()
+        rgbs = rgbs.reshape(B, T, 3, self.interp_shape[0], self.interp_shape[1])
+
+        device = rgbs.device
 
         queries[:, :, 1] *= self.interp_shape[1] / W
         queries[:, :, 2] *= self.interp_shape[0] / H
 
         if self.single_point:
-            traj_e = torch.zeros((B, T, N, 2)).cuda()
-            vis_e = torch.zeros((B, T, N)).cuda()
+            traj_e = torch.zeros((B, T, N, 2), device=device)
+            vis_e = torch.zeros((B, T, N), device=device)
             for pind in range((N)):
                 query = queries[:, pind : pind + 1]
 
@@ -60,8 +61,10 @@ class EvaluationPredictor(torch.nn.Module):
                 vis_e[:, t:, pind : pind + 1] = vis_e_pind[:, :, :1]
         else:
             if self.grid_size > 0:
-                xy = get_points_on_a_grid(self.grid_size, rgbs.shape[3:])
-                xy = torch.cat([torch.zeros_like(xy[:, :, :1]), xy], dim=2).cuda()  #
+                xy = get_points_on_a_grid(self.grid_size, rgbs.shape[3:], device=device)
+                xy = torch.cat([torch.zeros_like(xy[:, :, :1]), xy], dim=2).to(
+                    device
+                )  #
                 queries = torch.cat([queries, xy], dim=1)  #
 
             traj_e, __, vis_e, __ = self.model(
@@ -91,8 +94,8 @@ class EvaluationPredictor(torch.nn.Module):
             query = torch.cat([query, xy_target], dim=1).to(device)  #
 
         if self.grid_size > 0:
-            xy = get_points_on_a_grid(self.grid_size, rgbs.shape[3:])
-            xy = torch.cat([torch.zeros_like(xy[:, :, :1]), xy], dim=2).cuda()  #
+            xy = get_points_on_a_grid(self.grid_size, rgbs.shape[3:], device=device)
+            xy = torch.cat([torch.zeros_like(xy[:, :, :1]), xy], dim=2).to(device)  #
             query = torch.cat([query, xy], dim=1).to(device)  #
         # crop the video to start from the queried frame
         query[0, 0, 0] = 0

cotracker/predictor.py

@@ -116,7 +116,7 @@ class CoTrackerPredictor(torch.nn.Module):
             queries[:, :, 1] *= self.interp_shape[1] / W
             queries[:, :, 2] *= self.interp_shape[0] / H
         elif grid_size > 0:
-            grid_pts = get_points_on_a_grid(grid_size, self.interp_shape)
+            grid_pts = get_points_on_a_grid(grid_size, self.interp_shape, device=video.device)
             if segm_mask is not None:
                 segm_mask = F.interpolate(
                     segm_mask, tuple(self.interp_shape), mode="nearest"