up

app.py

@@ -4,23 +4,9 @@ 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}"
 
 
 @spaces.GPU
@@ -40,8 +26,8 @@ def yolov9_inference(img_path, model_id, image_size, conf_threshold, iou_thresho
     import yolov9
     
     # Load the model
-    model_path = attempt_download_from_hub(model_id)
-    model = yolov9.load(model_path, device="cpu")
+    model_path = download_models(model_id)
+    model = yolov9.load(model_path, device="cuda")
     
     # Set model parameters
     model.conf = conf_threshold
@@ -64,7 +50,10 @@ def app():
                 model_path = gr.Dropdown(
                     label="Model",
                     choices=[
-                        "kadirnar/yolov9-gelan-c",
+                        "gelan-c.pt",
+                        "gelan-e.pt",
+                        "yolov9-c.pt",
+                        "yolov9-e.pt",
                     ],
                     value="gelan-e.pt",
                 )
@@ -109,8 +98,8 @@ def app():
         gr.Examples(
             examples=[
                 [
-                    "data/zidane.jpg",
-                    "kadirnar/yolov9-gelan-c",
+                    "zidane.jpg"
+                    "gelan-e.pt",
                     640,
                     0.4,
                     0.5,