| from typing_extensions import Literal, TypeAlias |
|
|
|
|
| Processor_id: TypeAlias = Literal[ |
| "canny", "depth", "softedge", "lineart", "lineart_anime", "openpose", "normal", "tile", "none", "inpaint" |
| ] |
|
|
| class Annotator: |
| def __init__(self, processor_id: Processor_id, model_path="models/Annotators", detect_resolution=None, device='cuda', skip_processor=False): |
| if not skip_processor: |
| if processor_id == "canny": |
| from controlnet_aux.processor import CannyDetector |
| self.processor = CannyDetector() |
| elif processor_id == "depth": |
| from controlnet_aux.processor import MidasDetector |
| self.processor = MidasDetector.from_pretrained(model_path).to(device) |
| elif processor_id == "softedge": |
| from controlnet_aux.processor import HEDdetector |
| self.processor = HEDdetector.from_pretrained(model_path).to(device) |
| elif processor_id == "lineart": |
| from controlnet_aux.processor import LineartDetector |
| self.processor = LineartDetector.from_pretrained(model_path).to(device) |
| elif processor_id == "lineart_anime": |
| from controlnet_aux.processor import LineartAnimeDetector |
| self.processor = LineartAnimeDetector.from_pretrained(model_path).to(device) |
| elif processor_id == "openpose": |
| from controlnet_aux.processor import OpenposeDetector |
| self.processor = OpenposeDetector.from_pretrained(model_path).to(device) |
| elif processor_id == "normal": |
| from controlnet_aux.processor import NormalBaeDetector |
| self.processor = NormalBaeDetector.from_pretrained(model_path).to(device) |
| elif processor_id == "tile" or processor_id == "none" or processor_id == "inpaint": |
| self.processor = None |
| else: |
| raise ValueError(f"Unsupported processor_id: {processor_id}") |
| else: |
| self.processor = None |
|
|
| self.processor_id = processor_id |
| self.detect_resolution = detect_resolution |
| |
| def to(self,device): |
| if hasattr(self.processor,"model") and hasattr(self.processor.model,"to"): |
|
|
| self.processor.model.to(device) |
|
|
| def __call__(self, image, mask=None): |
| 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 |
|
|
|
|
