Removed cuda dependency

app.py

@@ -31,7 +31,7 @@ def install_add_dependencies():
         script = file.read()
     return call(script, shell=True)
 
-install_add_dependencies()
+#install_add_dependencies()
 
 class AppSteps(Enum):
     JUST_TEXT = 1
@@ -63,9 +63,6 @@ def get_args_parser():
 
     # training parameters
     parser.add_argument("--note", default="", help="add some notes to the experiment")
-    parser.add_argument(
-        "--device", default="cuda", help="device to use for training / testing"
-    )
     parser.add_argument("--resume", default="", help="resume from checkpoint")
     parser.add_argument(
         "--pretrain_model_path",
@@ -106,6 +103,10 @@ def get_args_parser():
 @spaces.GPU
 # Get counting model.
 def build_model_and_transforms(args):
+    if torch.cuda.is_available():
+        args.device = torch.device('cuda')
+    else:
+        args.device = torch.device('cpu')
     normalize = T.Compose(
         [T.ToTensor(), T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]
     )
@@ -208,25 +209,26 @@ def get_ind_to_filter(text, word_ids, keywords):
 
 @spaces.GPU
 def count(image, text, prompts, state):
+    global args
     print("state: " + str(state))
     keywords = "" # do not handle this for now
     # Handle no prompt case.
     if prompts is None:
         prompts = {"image": image, "points": []}
     input_image, _ = transform(image, {"exemplars": torch.tensor([])})
-    input_image = input_image.unsqueeze(0).cuda()
+    input_image = input_image.unsqueeze(0).to(args.device)
     exemplars = get_box_inputs(prompts["points"])
     print(exemplars)
     input_image_exemplars, exemplars = transform(prompts["image"], {"exemplars": torch.tensor(exemplars)})
-    input_image_exemplars = input_image_exemplars.unsqueeze(0).cuda()
-    exemplars = [exemplars["exemplars"].cuda()]
+    input_image_exemplars = input_image_exemplars.unsqueeze(0).to(args.device)
+    exemplars = [exemplars["exemplars"].to(args.device)]
     
     with torch.no_grad():
         model_output = model(
                 nested_tensor_from_tensor_list(input_image),
                 nested_tensor_from_tensor_list(input_image_exemplars),
                 exemplars,
-                [torch.tensor([0]).cuda() for _ in range(len(input_image))],
+                [torch.tensor([0]).to(args.device) for _ in range(len(input_image))],
                 captions=[text + " ."] * len(input_image),
             )
     
@@ -297,24 +299,25 @@ def count(image, text, prompts, state):
 
 @spaces.GPU
 def count_main(image, text, prompts):
+    global args
     keywords = "" # do not handle this for now
     # Handle no prompt case.
     if prompts is None:
         prompts = {"image": image, "points": []}
     input_image, _ = transform(image, {"exemplars": torch.tensor([])})
-    input_image = input_image.unsqueeze(0).cuda()
+    input_image = input_image.unsqueeze(0).to(args.device)
     exemplars = get_box_inputs(prompts["points"])
     print(exemplars)
     input_image_exemplars, exemplars = transform(prompts["image"], {"exemplars": torch.tensor(exemplars)})
-    input_image_exemplars = input_image_exemplars.unsqueeze(0).cuda()
-    exemplars = [exemplars["exemplars"].cuda()]
+    input_image_exemplars = input_image_exemplars.unsqueeze(0).to(args.device)
+    exemplars = [exemplars["exemplars"].to(args.device)]
     
     with torch.no_grad():
         model_output = model(
                 nested_tensor_from_tensor_list(input_image),
                 nested_tensor_from_tensor_list(input_image_exemplars),
                 exemplars,
-                [torch.tensor([0]).cuda() for _ in range(len(input_image))],
+                [torch.tensor([0]).to(args.device) for _ in range(len(input_image))],
                 captions=[text + " ."] * len(input_image),
             )