Upload app.py

app.py

@@ -41,7 +41,7 @@ def seed_everything(seed: int):
     os.environ["PYTHONHASHSEED"] = str(seed)
     np.random.seed(seed)
     torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)
+    # torch.cuda.manual_seed(seed)
     torch.backends.cudnn.deterministic = True
     torch.backends.cudnn.benchmark = False
 
@@ -55,15 +55,16 @@ model = Net(16)
 count_parameters(model)
 model.load_state_dict(state_dict['model'], strict=True)
 print('loaded ckpt')
-model.cuda()
+device = 'cpu:0'
+model.to(device)
 for n, p in model.named_parameters():
     p.requires_grad = False
 model.eval()
 
 tracker = Tracker(
     model=model,
-    mean=torch.tensor([0.485, 0.456, 0.406]).cuda().reshape(1, 3, 1, 1),
-    std=torch.tensor([0.229, 0.224, 0.225]).cuda().reshape(1, 3, 1, 1),
+    mean=torch.tensor([0.485, 0.456, 0.406]).to(device).reshape(1, 3, 1, 1),
+    std=torch.tensor([0.229, 0.224, 0.225]).to(device).reshape(1, 3, 1, 1),
     S=16,
     stride=8,
     inference_iters=4,
@@ -154,7 +155,8 @@ def process_video_with_points(video_path, click_points):
     frame_disps = []
     try:
         while True:
-            torch.cuda.empty_cache()
+            if 'cuda' in device:
+                torch.cuda.empty_cache()
             ret, frame = cap.read()
             if not ret:
                 break
@@ -195,7 +197,8 @@ def process_video_with_points(video_path, click_points):
     except RuntimeError as e:
         # Check if the error message indicates an OOM error.
         if "out of memory" in str(e).lower():
-            torch.cuda.empty_cache()
+            if 'cuda' in device:
+                torch.cuda.empty_cache()
             pbar.close()
             cap.release()
             print("Error: Out of Memory during video processing.")