diff --git a/CHANGELOG.md b/CHANGELOG.md index 1093ca8..d7b54c3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,12 @@ # Changelog +## 2024-05-20 + +- v24.5.1 +- uv를 사용하지 않게 함 +- 모든 허깅페이스 모델을 동시에 다운로드 시도함 +- 기본 탭 수를 2에서 4로 변경 + ## 2024-05-19 - v24.5.0 diff --git a/adetailer/__version__.py b/adetailer/__version__.py index 7ad6f1b..67b0a4b 100644 --- a/adetailer/__version__.py +++ b/adetailer/__version__.py @@ -1 +1 @@ -__version__ = "24.5.0" +__version__ = "24.5.1" diff --git a/adetailer/common.py b/adetailer/common.py index dfd7952..470bb2f 100644 --- a/adetailer/common.py +++ b/adetailer/common.py @@ -2,9 +2,10 @@ 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 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", diff --git a/adetailer/mask.py b/adetailer/mask.py index 3cd2edd..9496aa4 100644 --- a/adetailer/mask.py +++ b/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] diff --git a/adetailer/mediapipe.py b/adetailer/mediapipe.py index fd07619..25c6900 100644 --- a/adetailer/mediapipe.py +++ b/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() diff --git a/adetailer/ultralytics.py b/adetailer/ultralytics.py index 0c9fb86..dc93482 100644 --- a/adetailer/ultralytics.py +++ b/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) diff --git a/install.py b/install.py index 7a6d4d0..87fff5d 100644 --- a/install.py +++ b/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: diff --git a/scripts/!adetailer.py b/scripts/!adetailer.py index c5a5f16..0a02eb5 100644 --- a/scripts/!adetailer.py +++ b/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, ), )