Changed how device being allocated

app.py

@@ -57,6 +57,9 @@ def get_args_parser():
     >> python main.py --output_dir ./gdino_test -c config/cfg_fsc147_vit_b_test.py --eval --datasets config/datasets_fsc147.json --pretrain_model_path ../checkpoints_and_logs/gdino_train/checkpoint_best_regular.pth --options text_encoder_type=checkpoints/bert-base-uncased --sam_tt_norm --crop
     """
     parser = argparse.ArgumentParser("Set transformer detector", add_help=False)
+    parser.add_argument(
+        "--device", default="cuda", help="device to use for inference"
+    )
     parser.add_argument(
         "--options",
         nargs="+",
@@ -108,11 +111,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():
+        print("USING GPU")
         return torch.device('cuda')
     else:
+        print("USING CPU")
         return torch.device('cpu')
 
 # Get counting model.
@@ -136,8 +140,6 @@ def build_model_and_transforms(args):
         else:
             raise ValueError("Key {} can used by args only".format(k))
 
-    # Start with model on CPU.
-    args.device = "cpu"
     # fix the seed for reproducibility
     seed = 42
     torch.manual_seed(seed)
@@ -220,7 +222,6 @@ def get_ind_to_filter(text, word_ids, keywords):
 
 @spaces.GPU
 def count(image, text, prompts, state, device):
-    model.to(device)
     
     keywords = "" # do not handle this for now
     
@@ -303,13 +304,11 @@ def count(image, text, prompts, state, device):
             out_label += " " + str(exemplars[0].size()[0]) + " visual exemplars."
     else:
         out_label = "Nothing specified to detect."
-    #model.cpu()
+ 
     return (gr.Image(output_img, visible=True, label=out_label, show_label=True), gr.Number(label="Predicted Count", visible=True, value=boxes.shape[0]), new_submit_btn, gr.Tab(visible=True), step_3, state)
 
 @spaces.GPU
 def count_main(image, text, prompts, device):
-    model.to(device)
-    
     keywords = "" # do not handle this for now
     # Handle no prompt case.
     if prompts is None:
@@ -372,7 +371,7 @@ def count_main(image, text, prompts, device):
             out_label += " " + str(exemplars[0].size()[0]) + " visual exemplars."
     else:
         out_label = "Nothing specified to detect."
-    model.cpu()
+    
     return (gr.Image(output_img, visible=True, label=out_label, show_label=True), gr.Number(label="Predicted Count", visible=True, value=boxes.shape[0]))
 
 def remove_label(image):