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,39 +2,35 @@
 
 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
 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 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,
  ),
  )