feat: add ultralytics models, download concurrent

models.py

@@ -1,33 +1,40 @@
 from __future__ import annotations
 
+from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 
 from huggingface_hub import hf_hub_download
-from ultralytics import YOLO
 from tqdm.auto import tqdm
+from ultralytics import YOLO
 
-REPOBINGSU = "Bingsu/adetailer"
-REPOGUON = "guon/hand-eyes"
+BINGSU = "Bingsu/adetailer"
+GUON = "guon/hand-eyes"
+ULTRALYTICS = "Ultralytics/YOLOv8"
 
 model_info = [
-    (REPOBINGSU, "deepfashion2_yolov8s-seg.pt"),
-    (REPOBINGSU, "face_yolov8m.pt"),
-    (REPOBINGSU, "face_yolov8n.pt"),
-    (REPOBINGSU, "face_yolov8n_v2.pt"),
-    (REPOBINGSU, "face_yolov8s.pt"),
-    (REPOBINGSU, "face_yolov9c.pt"),
-    (REPOBINGSU, "hand_yolov8n.pt"),
-    (REPOBINGSU, "hand_yolov8s.pt"),
-    (REPOBINGSU, "hand_yolov9c.pt"),
-    (REPOBINGSU, "person_yolov8m-seg.pt"),
-    (REPOBINGSU, "person_yolov8n-seg.pt"),
-    (REPOBINGSU, "person_yolov8s-seg.pt"),
-    (REPOGUON, "Eyes.pt"),
-    (REPOGUON, "PitEyeDetailer-v1-seg.pt"),
-    (REPOGUON, "PitHandDetailer-v1-seg.pt"),
-    (REPOGUON, "female_breast-v4.2.pt"),
-    (REPOGUON, "full_eyes_detect_v1.pt"),
-    (REPOGUON, "lips_v1.pt"),
+    (BINGSU, "deepfashion2_yolov8s-seg.pt"),
+    (BINGSU, "face_yolov8m.pt"),
+    (BINGSU, "face_yolov8n.pt"),
+    (BINGSU, "face_yolov8n_v2.pt"),
+    (BINGSU, "face_yolov8s.pt"),
+    (BINGSU, "face_yolov9c.pt"),
+    (BINGSU, "hand_yolov8n.pt"),
+    (BINGSU, "hand_yolov8s.pt"),
+    (BINGSU, "hand_yolov9c.pt"),
+    (BINGSU, "person_yolov8m-seg.pt"),
+    (BINGSU, "person_yolov8n-seg.pt"),
+    (BINGSU, "person_yolov8s-seg.pt"),
+    (GUON, "Eyes.pt"),
+    (GUON, "PitEyeDetailer-v1-seg.pt"),
+    (GUON, "PitHandDetailer-v1-seg.pt"),
+    (GUON, "female_breast-v4.2.pt"),
+    (GUON, "full_eyes_detect_v1.pt"),
+    (GUON, "lips_v1.pt"),
+    (ULTRALYTICS, "yolov8l.pt"),
+    (ULTRALYTICS, "yolov8m.pt"),
+    (ULTRALYTICS, "yolov8n.pt"),
+    (ULTRALYTICS, "yolov8s.pt"),
+    (ULTRALYTICS, "yolov8x.pt"),
 ]
 
 
@@ -45,8 +52,17 @@ def export(path: str) -> str:
 def main():
     save = Path("models")
     save.mkdir(exist_ok=True)
-    for repo, model in tqdm(model_info):
-        path = download(model, repo)
+
+    pbar = tqdm(model_info)
+    with ThreadPoolExecutor() as executor:
+        fs = []
+        for repo, model in tqdm(model_info):
+            fu = executor.submit(download, model, repo)
+            fu.add_done_callback(lambda _: pbar.update())
+            fs.append(fu)
+
+    for future in fs:
+        path = future.result()
         onnx = export(path)
         target_dir = save / Path(model).stem
         (target_dir / "1").mkdir(exist_ok=True, parents=True)