diff --git a/modules/control/proc/teed/__init__.py b/modules/control/proc/teed/__init__.py new file mode 100644 index 000000000..2f0ab310b --- /dev/null +++ b/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