update

app.py

@@ -4,6 +4,24 @@ import os
 from huggingface_hub import hf_hub_download
 
 
+def attempt_download_from_hub(repo_id, hf_token=None):
+    # https://github.com/fcakyon/yolov5-pip/blob/main/yolov5/utils/downloads.py
+    from huggingface_hub import hf_hub_download, list_repo_files
+    from huggingface_hub.utils._errors import RepositoryNotFoundError
+    from huggingface_hub.utils._validators import HFValidationError
+    try:
+        repo_files = list_repo_files(repo_id=repo_id, repo_type='model', token=hf_token)
+        model_file = [f for f in repo_files if f.endswith('.pt')][0]
+        file = hf_hub_download(
+            repo_id=repo_id,
+            filename=model_file,
+            repo_type='model',
+            token=hf_token,
+        )
+        return file
+    except (RepositoryNotFoundError, HFValidationError):
+        return None
+    
 def download_models(model_id):
     hf_hub_download("merve/yolov9", filename=f"{model_id}", local_dir=f"./{model_id}")
     return f"./{model_id}"
@@ -26,7 +44,7 @@ def yolov9_inference(img_path, model_id, image_size, conf_threshold, iou_thresho
     import yolov9
     
     # Load the model
-    model_path = download_models(model_id)
+    model_path = attempt_download_from_hub(model_id, hf_token=None)
     model = yolov9.load(model_path, device="cuda")
     
     # Set model parameters