| 6 | ] |
| 7 | |
| 8 | class 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 CannyDetector |
| 13 | self.processor = CannyDetector() |
| 14 | elif processor_id == "depth": |
| 15 | from controlnet_aux.processor import MidasDetector |
| 16 | self.processor = MidasDetector.from_pretrained(model_path).to(device) |
| 17 | elif processor_id == "softedge": |
| 18 | from controlnet_aux.processor import HEDdetector |
| 19 | self.processor = HEDdetector.from_pretrained(model_path).to(device) |
| 20 | elif processor_id == "lineart": |
| 21 | from controlnet_aux.processor import LineartDetector |
| 22 | self.processor = LineartDetector.from_pretrained(model_path).to(device) |
| 23 | elif processor_id == "lineart_anime": |
| 24 | from controlnet_aux.processor import LineartAnimeDetector |
| 25 | self.processor = LineartAnimeDetector.from_pretrained(model_path).to(device) |
| 26 | elif processor_id == "openpose": |
| 27 | from controlnet_aux.processor import OpenposeDetector |
| 28 | self.processor = OpenposeDetector.from_pretrained(model_path).to(device) |
| 29 | elif processor_id == "normal": |
| 30 | from controlnet_aux.processor import NormalBaeDetector |
| 31 | 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 = None |
| 34 | else: |
| 35 | raise ValueError(f"Unsupported processor_id: {processor_id}") |
| 36 | else: |
| 37 | self.processor = None |
| 38 | |
| 39 | self.processor_id = processor_id |
| 40 | self.detect_resolution = detect_resolution |
| 41 | |
| 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.size |
| 49 | if self.processor_id == "openpose": |
| 50 | kwargs = { |
| 51 | "include_body": True, |
| 52 | "include_hand": True, |
| 53 | "include_face": True |
| 54 | } |
| 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 image |
| 62 |
no outgoing calls
no test coverage detected