Update app.py

app.py

@@ -24,8 +24,29 @@ category_dict = {
     77: 'teddy bear', 78: 'hair drier', 79: 'toothbrush'
 }
 
+
+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
+
+
 @spaces.GPU(duration=200)
 def LeYOLO_inference(image, model_id, image_size, conf_threshold, iou_threshold):
+    MODEL_PATH = attempt_download_from_hub("kadirnar/yolov10x")
     model = model = YOLO(f"kadirnar/{model_id}")
     results = model(source=image, imgsz=image_size, iou=iou_threshold, conf=conf_threshold, verbose=False)[0]
     detections = sv.Detections.from_ultralytics(results)
@@ -48,12 +69,12 @@ def app():
                 model_id = gr.Dropdown(
                     label="Model",
                     choices=[
-                        "LeYOLOSmall",
-                        "LeYOLONano",
-                        "LeYOLOMedium",
-                        "LeYOLOLarge",
+                        "kadirnar/LeYOLOSmall",
+                        "kadirnar/LeYOLONano",
+                        "kadirnar/LeYOLOMedium",
+                        "kadirnar/LeYOLOLarge",
                     ],
-                    value="LeYOLOMedium",
+                    value="kadirnar/LeYOLOMedium",
                 )
                 image_size = gr.Slider(
                     label="Image Size",
@@ -97,14 +118,14 @@ def app():
             examples=[
                 [
                     "dog.jpeg",
-                    "LeYOLONano",
+                    "kadirnar/LeYOLOMedium",
                     640,
                     0.25,
                     0.45,
                 ],
                 [
                     "zidane.jpg",
-                    "LeYOLOMedium",
+                    "kadirnar/LeYOLOMedium",
                     640,
                     0.25,
                     0.45,