bug fix

app.py

@@ -33,13 +33,13 @@ ckpt_repo_id = "ShilongLiu/GroundingDINO"
 ckpt_filenmae = "groundingdino_swinb_cogcoor.pth"
 
 
-def load_model_hf(model_config_path, repo_id, filename, device='cpu'):
+def load_model_hf(model_config_path, repo_id, filename, device='cuda'):
     args = SLConfig.fromfile(model_config_path) 
     model = build_model(args)
     args.device = device
 
     cache_file = hf_hub_download(repo_id=repo_id, filename=filename)
-    checkpoint = torch.load(cache_file, map_location='cpu')
+    checkpoint = torch.load(cache_file, map_location=device)
     log = model.load_state_dict(clean_state_dict(checkpoint['model']), strict=False)
     print("Model loaded from {} \n => {}".format(cache_file, log))
     _ = model.eval()
@@ -72,7 +72,7 @@ def run_grounding(input_image, grounding_caption, box_threshold, text_threshold)
     image_pil: Image = image_transform_grounding_for_vis(init_image)
 
     # run grounidng
-    boxes, logits, phrases = predict(model, image_tensor, grounding_caption, box_threshold, text_threshold, device='cpu')
+    boxes, logits, phrases = predict(model, image_tensor, grounding_caption, box_threshold, text_threshold, device='cuda')
     annotated_frame = annotate(image_source=np.asarray(image_pil), boxes=boxes, logits=logits, phrases=phrases)
     image_with_box = Image.fromarray(cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB))
 
@@ -96,7 +96,7 @@ if __name__ == "__main__":
     with gr.Blocks(css=css) as demo:
         gr.Markdown("<h1><center>Grounding DINO<h1><center>")
         gr.Markdown("<h3><center>Open-World Detection with <a href='https://github.com/IDEA-Research/GroundingDINO'>Grounding DINO</a><h3><center>")
-        gr.Markdown("<h3><center>Note the model runs on CPU, so it may take a while to run the model.<h3><center>")
+        gr.Markdown("<h3><center>Running on GPU for faster inference<h3><center>")
 
         with gr.Row():
             with gr.Column():