Merge branch 'dev'

CHANGELOG.md

@@ -1,5 +1,12 @@
 # Changelog
 
+## 2024-05-20
+
+- v24.5.1
+- uv를 사용하지 않게 함
+- 모든 허깅페이스 모델을 동시에 다운로드 시도함
+- 기본 탭 수를 2에서 4로 변경
+
 ## 2024-05-19
 
 - v24.5.0

adetailer/__version__.py

@@ -1 +1 @@
-__version__ = "24.5.0"
+__version__ = "24.5.1"

adetailer/common.py

@@ -2,9 +2,10 @@ from __future__ import annotations
 
 import os
 from collections import OrderedDict
+from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any, Generic, Optional, TypeVar
 
 from huggingface_hub import hf_hub_download
 from PIL import Image, ImageDraw
@@ -12,29 +13,24 @@ from rich import print
 from torchvision.transforms.functional import to_pil_image
 
 REPO_ID = "Bingsu/adetailer"
-_download_failed = False
+
+T = TypeVar("T", int, float)
 
 
 @dataclass
-class PredictOutput:
-    bboxes: list[list[int | float]] = field(default_factory=list)
+class PredictOutput(Generic[T]):
+    bboxes: list[list[T]] = field(default_factory=list)
     masks: list[Image.Image] = field(default_factory=list)
     preview: Optional[Image.Image] = None
 
 
-def hf_download(file: str, repo_id: str = REPO_ID) -> str | None:
-    global _download_failed
-
-    if _download_failed:
-        return "INVALID"
-
+def hf_download(file: str, repo_id: str = REPO_ID) -> str:
     try:
         path = hf_hub_download(repo_id, file)
     except Exception:
         msg = f"[-] ADetailer: Failed to load model {file!r} from huggingface"
         print(msg)
         path = "INVALID"
-        _download_failed = True
     return path
 
 
@@ -50,6 +46,19 @@ def scan_model_dir(path: Path) -> list[Path]:
     return [p for p in path.rglob("*") if p.is_file() and p.suffix == ".pt"]
 
 
+def download_models(*names: str) -> dict[str, str]:
+    models = OrderedDict()
+    with ThreadPoolExecutor() as executor:
+        for name in names:
+            if "-world" in name:
+                models[name] = executor.submit(
+                    hf_download, name, repo_id="Bingsu/yolo-world-mirror"
+                )
+            else:
+                models[name] = executor.submit(hf_download, name)
+    return {name: future.result() for name, future in models.items()}
+
+
 def get_models(
     *dirs: str | os.PathLike[str], huggingface: bool = True
 ) -> OrderedDict[str, str]:
@@ -62,18 +71,16 @@ def get_models(
 
     models = OrderedDict()
     if huggingface:
-        models.update(
-            {
-                "face_yolov8n.pt": hf_download("face_yolov8n.pt"),
-                "face_yolov8s.pt": hf_download("face_yolov8s.pt"),
-                "hand_yolov8n.pt": hf_download("hand_yolov8n.pt"),
-                "person_yolov8n-seg.pt": hf_download("person_yolov8n-seg.pt"),
-                "person_yolov8s-seg.pt": hf_download("person_yolov8s-seg.pt"),
-                "yolov8x-worldv2.pt": hf_download(
-                    "yolov8x-worldv2.pt", repo_id="Bingsu/yolo-world-mirror"
-                ),
-            }
-        )
+        to_download = [
+            "face_yolov8n.pt",
+            "face_yolov8s.pt",
+            "hand_yolov8n.pt",
+            "person_yolov8n-seg.pt",
+            "person_yolov8s-seg.pt",
+            "yolov8x-worldv2.pt",
+        ]
+        models.update(download_models(*to_download))
+
     models.update(
         {
             "mediapipe_face_full": "mediapipe_face_full",

adetailer/mask.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 from enum import IntEnum
 from functools import partial, reduce
 from math import dist
-from typing import Any
+from typing import Any, TypeVar
 
 import cv2
 import numpy as np
@@ -26,6 +26,9 @@ class MergeInvert(IntEnum):
     MERGE_INVERT = 2
 
 
+T = TypeVar("T", int, float)
+
+
 def _dilate(arr: np.ndarray, value: int) -> np.ndarray:
     kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (value, value))
     return cv2.dilate(arr, kernel, iterations=1)
@@ -96,7 +99,7 @@ def has_intersection(im1: Any, im2: Any) -> bool:
     return not is_all_black(cv2.bitwise_and(arr1, arr2))
 
 
-def bbox_area(bbox: list[float]) -> float:
+def bbox_area(bbox: list[T]) -> T:
     return (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
 
 
@@ -141,25 +144,25 @@ def mask_preprocess(
 
 
 # Bbox sorting
-def _key_left_to_right(bbox: list[float]) -> float:
+def _key_left_to_right(bbox: list[T]) -> T:
     """
     Left to right
 
     Parameters
     ----------
-    bbox: list[float]
+    bbox: list[int] | list[float]
         list of [x1, y1, x2, y2]
     """
     return bbox[0]
 
 
-def _key_center_to_edge(bbox: list[float], *, center: tuple[float, float]) -> float:
+def _key_center_to_edge(bbox: list[T], *, center: tuple[float, float]) -> float:
     """
     Center to edge
 
     Parameters
     ----------
-    bbox: list[float]
+    bbox: list[int] | list[float]
         list of [x1, y1, x2, y2]
     image: Image.Image
         the image
@@ -168,21 +171,21 @@ def _key_center_to_edge(bbox: list[float], *, center: tuple[float, float]) -> fl
     return dist(center, bbox_center)
 
 
-def _key_area(bbox: list[float]) -> float:
+def _key_area(bbox: list[T]) -> T:
     """
     Large to small
 
     Parameters
     ----------
-    bbox: list[float]
+    bbox: list[int] | list[float]
         list of [x1, y1, x2, y2]
     """
     return -bbox_area(bbox)
 
 
 def sort_bboxes(
-    pred: PredictOutput, order: int | SortBy = SortBy.NONE
-) -> PredictOutput:
+    pred: PredictOutput[T], order: int | SortBy = SortBy.NONE
+) -> PredictOutput[T]:
     if order == SortBy.NONE or len(pred.bboxes) <= 1:
         return pred
 
@@ -205,12 +208,14 @@ def sort_bboxes(
 
 
 # Filter by ratio
-def is_in_ratio(bbox: list[float], low: float, high: float, orig_area: int) -> bool:
+def is_in_ratio(bbox: list[T], low: float, high: float, orig_area: int) -> bool:
     area = bbox_area(bbox)
     return low <= area / orig_area <= high
 
 
-def filter_by_ratio(pred: PredictOutput, low: float, high: float) -> PredictOutput:
+def filter_by_ratio(
+    pred: PredictOutput[T], low: float, high: float
+) -> PredictOutput[T]:
     if not pred.bboxes:
         return pred
 
@@ -223,7 +228,7 @@ def filter_by_ratio(pred: PredictOutput, low: float, high: float) -> PredictOutp
     return pred
 
 
-def filter_k_largest(pred: PredictOutput, k: int = 0) -> PredictOutput:
+def filter_k_largest(pred: PredictOutput[T], k: int = 0) -> PredictOutput[T]:
     if not pred.bboxes or k == 0:
         return pred
     areas = [bbox_area(bbox) for bbox in pred.bboxes]

adetailer/mediapipe.py

@@ -28,7 +28,7 @@ def mediapipe_predict(
 
 def mediapipe_face_detection(
     model_type: int, image: Image.Image, confidence: float = 0.3
-) -> PredictOutput:
+) -> PredictOutput[float]:
     import mediapipe as mp
 
     img_width, img_height = image.size
@@ -68,7 +68,9 @@ def mediapipe_face_detection(
     return PredictOutput(bboxes=bboxes, masks=masks, preview=preview)
 
 
-def mediapipe_face_mesh(image: Image.Image, confidence: float = 0.3) -> PredictOutput:
+def mediapipe_face_mesh(
+    image: Image.Image, confidence: float = 0.3
+) -> PredictOutput[int]:
     import mediapipe as mp
 
     mp_face_mesh = mp.solutions.face_mesh
@@ -98,7 +100,9 @@ def mediapipe_face_mesh(image: Image.Image, confidence: float = 0.3) -> PredictO
                 connection_drawing_spec=drawing_styles.get_default_face_mesh_tesselation_style(),
             )
 
-            points = np.intp([(land.x * w, land.y * h) for land in landmarks.landmark])
+            points = np.array(
+                [[land.x * w, land.y * h] for land in landmarks.landmark], dtype=int
+            )
             outline = cv2.convexHull(points).reshape(-1).tolist()
 
             mask = Image.new("L", image.size, "black")
@@ -113,7 +117,7 @@ def mediapipe_face_mesh(image: Image.Image, confidence: float = 0.3) -> PredictO
 
 def mediapipe_face_mesh_eyes_only(
     image: Image.Image, confidence: float = 0.3
-) -> PredictOutput:
+) -> PredictOutput[int]:
     import mediapipe as mp
 
     mp_face_mesh = mp.solutions.face_mesh
@@ -136,7 +140,9 @@ def mediapipe_face_mesh_eyes_only(
         masks = []
 
         for landmarks in pred.multi_face_landmarks:
-            points = np.intp([(land.x * w, land.y * h) for land in landmarks.landmark])
+            points = np.array(
+                [[land.x * w, land.y * h] for land in landmarks.landmark], dtype=int
+            )
             left_eyes = points[left_idx]
             right_eyes = points[right_idx]
             left_outline = cv2.convexHull(left_eyes).reshape(-1).tolist()

adetailer/ultralytics.py

@@ -21,7 +21,7 @@ def ultralytics_predict(
     confidence: float = 0.3,
     device: str = "",
     classes: str = "",
-) -> PredictOutput:
+) -> PredictOutput[float]:
     from ultralytics import YOLO
 
     model = YOLO(model_path)

install.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import importlib.util
-import os
 import subprocess
 import sys
 from importlib.metadata import version  # python >= 3.8
@@ -39,11 +38,7 @@ def is_installed(
 
 
 def run_pip(*args):
-    subprocess.run([sys.executable, "-m", "pip", "install", *args], check=False)
-
-
-def run_uv_pip(*args):
-    subprocess.run([sys.executable, "-m", "uv", "pip", "install", *args], check=False)
+    subprocess.run([sys.executable, "-m", "pip", "install", *args], check=True)
 
 
 def install():
@@ -56,11 +51,6 @@ def install():
         ("protobuf", "4.25.3", "4.9999"),
     ]
 
-    if not is_installed("uv", "0.1.44", None):
-        run_pip("uv>=0.1.44")
-
-    os.environ["UV_PYTHON"] = sys.executable
-
     pkgs = []
     for pkg, low, high in deps:
         if not is_installed(pkg, low, high):
@@ -74,7 +64,8 @@ def install():
                 cmd = pkg
             pkgs.append(cmd)
 
-    run_uv_pip(*pkgs)
+    if pkgs:
+        run_pip(*pkgs)
 
 
 try:

scripts/!adetailer.py

@@ -853,10 +853,10 @@ def on_ui_settings():
     shared.opts.add_option(
         "ad_max_models",
         shared.OptionInfo(
-            default=2,
+            default=4,
             label="Max models",
             component=gr.Slider,
-            component_args={"minimum": 1, "maximum": 10, "step": 1},
+            component_args={"minimum": 1, "maximum": 15, "step": 1},
             section=section,
         ),
     )