| from typing import List |
|
|
| import cv2 |
| import numpy as np |
| import torch |
| from controlnet_aux import OpenposeDetector |
| from diffusers import (ControlNetModel, StableDiffusionControlNetPipeline, |
| UniPCMultistepScheduler) |
| from PIL import Image |
| from util.cache import clear_cuda_and_gc |
| from util.commons import disable_safety_checker, download_image |
|
|
|
|
| class ControlNet: |
| __current_task_name = "" |
|
|
| def load(self, model_dir: str): |
| |
| self.load_canny() |
|
|
| pipe = StableDiffusionControlNetPipeline.from_pretrained( |
| model_dir, controlnet=self.controlnet, torch_dtype=torch.float16 |
| ) |
| pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) |
| pipe.enable_model_cpu_offload() |
| pipe.enable_xformers_memory_efficient_attention() |
| disable_safety_checker(pipe) |
| self.pipe = pipe |
|
|
| def load_canny(self): |
| if self.__current_task_name == "canny": |
| return |
| canny = ControlNetModel.from_pretrained( |
| "lllyasviel/control_v11p_sd15_canny", torch_dtype=torch.float16 |
| ).to("cuda") |
| self.__current_task_name = "canny" |
| self.controlnet = canny |
| if hasattr(self, "pipe"): |
| self.pipe.controlnet = canny |
| clear_cuda_and_gc() |
|
|
| def load_pose(self): |
| if self.__current_task_name == "pose": |
| return |
| pose = ControlNetModel.from_pretrained( |
| "lllyasviel/control_v11p_sd15_openpose", torch_dtype=torch.float16 |
| ).to("cuda") |
| self.__current_task_name = "pose" |
| self.controlnet = pose |
| if hasattr(self, "pipe"): |
| self.pipe.controlnet = pose |
| clear_cuda_and_gc() |
|
|
| def cleanup(self): |
| self.pipe.controlnet = None |
| self.controlnet = None |
| self.__current_task_name = "" |
|
|
| clear_cuda_and_gc() |
|
|
| @torch.inference_mode() |
| def process_canny( |
| self, |
| prompt: List[str], |
| imageUrl: str, |
| seed: int, |
| steps: int, |
| negative_prompt: List[str], |
| height: int, |
| width: int, |
| ): |
| if self.__current_task_name != "canny": |
| raise Exception("ControlNet is not loaded with canny model") |
|
|
| torch.manual_seed(seed) |
|
|
| init_image = download_image(imageUrl) |
| init_image = self.__canny_detect_edge(init_image) |
|
|
| return self.pipe.__call__( |
| prompt=prompt, |
| image=init_image, |
| guidance_scale=9, |
| num_images_per_prompt=1, |
| negative_prompt=negative_prompt, |
| num_inference_steps=steps, |
| height=height, |
| width=width, |
| ).images |
|
|
| @torch.inference_mode() |
| def process_pose( |
| self, |
| prompt: List[str], |
| image: List[Image.Image], |
| seed: int, |
| steps: int, |
| negative_prompt: List[str], |
| height: int, |
| width: int, |
| ): |
| if self.__current_task_name != "pose": |
| raise Exception("ControlNet is not loaded with pose model") |
|
|
| torch.manual_seed(seed) |
|
|
| return self.pipe.__call__( |
| prompt=prompt, |
| image=image, |
| num_images_per_prompt=1, |
| num_inference_steps=steps, |
| negative_prompt=negative_prompt, |
| height=height, |
| width=width, |
| ).images |
|
|
| def detect_pose(self, imageUrl: str) -> Image.Image: |
| detector = OpenposeDetector.from_pretrained("lllyasviel/ControlNet") |
| image = download_image(imageUrl) |
| image = detector.__call__(image) |
| return image |
|
|
| def __canny_detect_edge(self, image: Image.Image) -> Image.Image: |
| image_array = np.array(image) |
|
|
| low_threshold = 100 |
| high_threshold = 200 |
|
|
| image_array = cv2.Canny(image_array, low_threshold, high_threshold) |
| image_array = image_array[:, :, None] |
| image_array = np.concatenate([image_array, image_array, image_array], axis=2) |
| canny_image = Image.fromarray(image_array) |
| return canny_image |
|
|
