| from typing_extensions import Literal, TypeAlias | |
| import warnings | |
| with warnings.catch_warnings(): | |
| warnings.simplefilter("ignore") | |
| from controlnet_aux.processor import ( | |
| CannyDetector, MidasDetector, HEDdetector, LineartDetector, LineartAnimeDetector, OpenposeDetector | |
| ) | |
| Processor_id: TypeAlias = Literal[ | |
| "canny", "depth", "softedge", "lineart", "lineart_anime", "openpose", "tile" | |
| ] | |
| class Annotator: | |
| def __init__(self, processor_id: Processor_id, model_path="models/Annotators", detect_resolution=None): | |
| if processor_id == "canny": | |
| self.processor = CannyDetector() | |
| elif processor_id == "depth": | |
| self.processor = MidasDetector.from_pretrained(model_path).to("cuda") | |
| elif processor_id == "softedge": | |
| self.processor = HEDdetector.from_pretrained(model_path).to("cuda") | |
| elif processor_id == "lineart": | |
| self.processor = LineartDetector.from_pretrained(model_path).to("cuda") | |
| elif processor_id == "lineart_anime": | |
| self.processor = LineartAnimeDetector.from_pretrained(model_path).to("cuda") | |
| elif processor_id == "openpose": | |
| self.processor = OpenposeDetector.from_pretrained(model_path).to("cuda") | |
| elif processor_id == "tile": | |
| self.processor = None | |
| else: | |
| raise ValueError(f"Unsupported processor_id: {processor_id}") | |
| self.processor_id = processor_id | |
| self.detect_resolution = detect_resolution | |
| def __call__(self, image): | |
| width, height = image.size | |
| if self.processor_id == "openpose": | |
| kwargs = { | |
| "include_body": True, | |
| "include_hand": True, | |
| "include_face": True | |
| } | |
| else: | |
| kwargs = {} | |
| if self.processor is not None: | |
| detect_resolution = self.detect_resolution if self.detect_resolution is not None else min(width, height) | |
| image = self.processor(image, detect_resolution=detect_resolution, image_resolution=min(width, height), **kwargs) | |
| image = image.resize((width, height)) | |
| return image | |