merge: modules/control/proc/teed/__init__.py

modules/control/proc/teed/__init__.py

@@ -0,0 +1,21 @@
+import numpy as np
+from PIL import Image
+
+
+class TEEDDetector:
+    def __init__(self, model):
+        self.model = model
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_or_path="fal/teed", cache_dir=None, local_files_only=False):
+        from installer import install
+        install('controlnet-aux', quiet=True)
+        from controlnet_aux import TEEDdetector as _TEEDdetector
+        model = _TEEDdetector.from_pretrained(pretrained_model_or_path, filename="5_model.pth", cache_dir=cache_dir)
+        return cls(model)
+
+    def __call__(self, image, output_type="pil", **kwargs):
+        if isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
+        result = self.model(image, output_type=output_type)
+        return result