| | 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 |
| |
|
