Added @spaces.gpu decorator and switched to gpu officially

app.py

@@ -28,34 +28,18 @@ cwd = os.getcwd()
 print("Current working directory:", cwd)
 
 # Installing dependencies not in requirements.txt
-@spaces.GPU
-def install_add_dependencies():
-    print("inside install_add_dependencies")
-    print(torch.cuda.is_available())
-    with open('./build_ops.sh', 'rb') as file:
-        script = file.read()
-    return call(script, shell=True)
-
-def build_custom_prompter():
-    with open('./build_custom_prompter.sh', 'rb') as file:
-        script = file.read()
-    return call(script, shell=True)
-
-def build_multiscale_deform():
-    with open('./build_multiscale_deform.sh', 'rb') as file:
-        script = file.read()
-    return call(script, shell=True)
-
-build_custom_prompter()
+subprocess.run(
+    shlex.split(
+        "pip install gradio_image_prompter-0.1.0-py3-none-any.whl"
+    )
+)
 from gradio_image_prompter import ImagePrompter
+
 subprocess.run(
     shlex.split(
         "pip install MultiScaleDeformableAttention-1.0-cp310-cp310-linux_x86_64.whl"
     )
 )
-#print("torch version")
-#print(torch.version.cuda)
-#install_add_dependencies()
 
 class AppSteps(Enum):
     JUST_TEXT = 1
@@ -124,6 +108,12 @@ def get_args_parser():
     parser.add_argument("--amp", action="store_true", help="Train with mixed precision")
     return parser
 
+@spaces.GPU
+def get_device():
+    if torch.cuda.is_available():
+        return torch.device('cuda')
+    else:
+        return torch.device('cpu')
 
 # Get counting model.
 @spaces.GPU
@@ -162,8 +152,6 @@ def build_model_and_transforms(args):
     build_func = MODULE_BUILD_FUNCS.get(args.modelname)
     model, _, _ = build_func(args)
 
-    #model.to(device)
-
     checkpoint = torch.load(args.pretrain_model_path, map_location="cpu")["model"]
     model.load_state_dict(checkpoint, strict=False)
 
@@ -174,11 +162,8 @@ def build_model_and_transforms(args):
 
 parser = argparse.ArgumentParser("Counting Application", parents=[get_args_parser()])
 args = parser.parse_args()
-#if torch.cuda.is_available():
-#    args.device = torch.device('cuda')
-#else:
-#    args.device = torch.device('cpu')
-args.device = torch.device('cpu')
+
+args.device = get_device()
 model, transform = build_model_and_transforms(args)
 
 examples = [
@@ -233,11 +218,12 @@ def get_ind_to_filter(text, word_ids, keywords):
 
     return inds_to_filter
 
-#@spaces.GPU
+@spaces.GPU
 def count(image, text, prompts, state, device):
     model.to(device)
-    print("state: " + str(state))
+    
     keywords = "" # do not handle this for now
+    
     # Handle no prompt case.
     if prompts is None:
         prompts = {"image": image, "points": []}
@@ -259,11 +245,7 @@ def count(image, text, prompts, state, device):
             )
     
     ind_to_filter = get_ind_to_filter(text, model_output["token"][0].word_ids, keywords)
-    print(model_output["token"][0].tokens)
-    print(ind_to_filter)
-    print(model_output["pred_logits"].sigmoid()[0].shape)
     logits = model_output["pred_logits"].sigmoid()[0][:, ind_to_filter]
-    print(logits.shape)
     boxes = model_output["pred_boxes"][0]
     if len(keywords.strip()) > 0:
         box_mask = (logits > CONF_THRESH).sum(dim=-1) == len(ind_to_filter)
@@ -339,7 +321,6 @@ def count_main(image, text, prompts, device):
     input_image_exemplars, exemplars = transform(prompts["image"], {"exemplars": torch.tensor(exemplars)})
     input_image_exemplars = input_image_exemplars.unsqueeze(0).to(device)
     exemplars = [exemplars["exemplars"].to(device)]
-    print("image device: " + str(input_image.device))
     
     with torch.no_grad():
         model_output = model(
@@ -351,11 +332,7 @@ def count_main(image, text, prompts, device):
             )
     
     ind_to_filter = get_ind_to_filter(text, model_output["token"][0].word_ids, keywords)
-    print(model_output["token"][0].tokens)
-    print(ind_to_filter)
-    print(model_output["pred_logits"].sigmoid()[0].shape)
     logits = model_output["pred_logits"].sigmoid()[0][:, ind_to_filter]
-    print(logits.shape)
     boxes = model_output["pred_boxes"][0]
     if len(keywords.strip()) > 0:
         box_mask = (logits > CONF_THRESH).sum(dim=-1) == len(ind_to_filter)