hugging-apps/echo-memory

1from typing_extensions import Literal, TypeAlias2 3 4Processor_id: TypeAlias = Literal[5    "canny", "depth", "softedge", "lineart", "lineart_anime", "openpose", "normal", "tile", "none", "inpaint"6]7 8class Annotator:9    def __init__(self, processor_id: Processor_id, model_path="models/Annotators", detect_resolution=None, device='cuda', skip_processor=False):10        if not skip_processor:11            if processor_id == "canny":12                from controlnet_aux.processor import CannyDetector13                self.processor = CannyDetector()14            elif processor_id == "depth":15                from controlnet_aux.processor import MidasDetector16                self.processor = MidasDetector.from_pretrained(model_path).to(device)17            elif processor_id == "softedge":18                from controlnet_aux.processor import HEDdetector19                self.processor = HEDdetector.from_pretrained(model_path).to(device)20            elif processor_id == "lineart":21                from controlnet_aux.processor import LineartDetector22                self.processor = LineartDetector.from_pretrained(model_path).to(device)23            elif processor_id == "lineart_anime":24                from controlnet_aux.processor import LineartAnimeDetector25                self.processor = LineartAnimeDetector.from_pretrained(model_path).to(device)26            elif processor_id == "openpose":27                from controlnet_aux.processor import OpenposeDetector28                self.processor = OpenposeDetector.from_pretrained(model_path).to(device)29            elif processor_id == "normal":30                from controlnet_aux.processor import NormalBaeDetector31                self.processor = NormalBaeDetector.from_pretrained(model_path).to(device)32            elif processor_id == "tile" or processor_id == "none" or processor_id == "inpaint":33                self.processor = None34            else:35                raise ValueError(f"Unsupported processor_id: {processor_id}")36        else:37            self.processor = None38 39        self.processor_id = processor_id40        self.detect_resolution = detect_resolution41    42    def to(self,device):43        if hasattr(self.processor,"model") and hasattr(self.processor.model,"to"):44 45            self.processor.model.to(device)46 47    def __call__(self, image, mask=None):48        width, height = image.size49        if self.processor_id == "openpose":50            kwargs = {51                "include_body": True,52                "include_hand": True,53                "include_face": True54            }55        else:56            kwargs = {}57        if self.processor is not None:58            detect_resolution = self.detect_resolution if self.detect_resolution is not None else min(width, height)59            image = self.processor(image, detect_resolution=detect_resolution, image_resolution=min(width, height), **kwargs)60        image = image.resize((width, height))61        return image62 63