Adding YOLO detectors for specific use cases

asdff/__init__.py

@@ -1,10 +1,20 @@
 from .__version__ import __version__
 from .sd import AdCnPipeline, AdPipeline
-from .yolo import yolo_detector
+from .yolo import (
+    yolo_detector,
+    yolo_face_detector,
+    yolo_hand_detector,
+    yolo_person_detector,
+    yolo_clothes_detector
+)
 
 __all__ = [
     "AdPipeline",
     "AdCnPipeline",
     "yolo_detector",
+    "yolo_face_detector",
+    "yolo_hand_detector",
+    "yolo_person_detector",
+    "yolo_clothes_detector",
     "__version__",
 ]

asdff/base.py

@@ -14,7 +14,7 @@
     mask_dilate,
     mask_gaussian_blur,
 )
-from asdff.yolo import yolo_detector
+from asdff.yolo import yolo_face_detector
 
 logger = logging.get_logger("diffusers")
 
@@ -122,7 +122,7 @@ def __call__(  # noqa: C901
 
     @property
     def default_detector(self) -> Callable[..., list[Image.Image] | None]:
-        return yolo_detector
+        return yolo_face_detector
 
     def _get_txt2img_args(
         self, common: Mapping[str, Any], txt2img_only: Mapping[str, Any]

asdff/yolo.py

@@ -78,3 +78,31 @@ def yolo_detector(
         masks = mask_to_pil(pred[0].masks.data, image.size)
 
     return masks
+
+def yolo_face_detector(
+        image: Image.Image,
+        confidence: float = 0.3
+    ) -> list[Image.Image] | None:
+    model_path = hf_hub_download("Bingsu/adetailer", "face_yolov8n.pt")
+    return yolo_detector(image, model_path=model_path, confidence=confidence)
+
+def yolo_hand_detector(
+        image: Image.Image,
+        confidence: float = 0.3
+    ) -> list[Image.Image] | None:
+    model_path = hf_hub_download("Bingsu/adetailer", "hand_yolov9c.pt")
+    return yolo_detector(image, model_path=model_path, confidence=confidence)
+
+def yolo_person_detector(
+        image: Image.Image,
+        confidence: float = 0.3
+    ) -> list[Image.Image] | None:
+    model_path = hf_hub_download("Bingsu/adetailer", "person_yolov8s-seg.pt")
+    return yolo_detector(image, model_path=model_path, confidence=confidence)
+
+def yolo_clothes_detector(
+        image: Image.Image,
+        confidence: float = 0.3
+    ) -> list[Image.Image] | None:
+    model_path = hf_hub_download("Bingsu/adetailer", "deepfashion2_yolov8s-seg.pt")
+    return yolo_detector(image, model_path=model_path, confidence=confidence)