hysts/ControlNet
993
1# This file is adapted from gradio_*.py in https://github.com/lllyasviel/ControlNet/tree/f4748e3630d8141d7765e2bd9b1e348f478477072# The original license file is LICENSE.ControlNet in this repo.3from __future__ import annotations4 5import gc6import pathlib7import sys8 9import cv210import numpy as np11import PIL.Image12import torch13from diffusers import (ControlNetModel, DiffusionPipeline,14 StableDiffusionControlNetPipeline,15 UniPCMultistepScheduler)16 17repo_dir = pathlib.Path(__file__).parent18submodule_dir = repo_dir / 'ControlNet'19sys.path.append(submodule_dir.as_posix())20 21try:22 from annotator.canny import apply_canny23 from annotator.hed import apply_hed, nms24 from annotator.midas import apply_midas25 from annotator.mlsd import apply_mlsd26 from annotator.openpose import apply_openpose27 from annotator.uniformer import apply_uniformer28 from annotator.util import HWC3, resize_image29except Exception:30 pass31 32CONTROLNET_MODEL_IDS = {33 'canny': 'lllyasviel/sd-controlnet-canny',34 'hough': 'lllyasviel/sd-controlnet-mlsd',35 'hed': 'lllyasviel/sd-controlnet-hed',36 'scribble': 'lllyasviel/sd-controlnet-scribble',37 'pose': 'lllyasviel/sd-controlnet-openpose',38 'seg': 'lllyasviel/sd-controlnet-seg',39 'depth': 'lllyasviel/sd-controlnet-depth',40 'normal': 'lllyasviel/sd-controlnet-normal',41}42 43 44def download_all_controlnet_weights() -> None:45 for model_id in CONTROLNET_MODEL_IDS.values():46 ControlNetModel.from_pretrained(model_id)47 48 49class Model:50 def __init__(self,51 base_model_id: str = 'runwayml/stable-diffusion-v1-5',52 task_name: str = 'canny'):53 self.device = torch.device(54 'cuda:0' if torch.cuda.is_available() else 'cpu')55 self.base_model_id = ''56 self.task_name = ''57 self.pipe = self.load_pipe(base_model_id, task_name)58 59 def load_pipe(self, base_model_id: str, task_name) -> DiffusionPipeline:60 if self.device.type == 'cpu':61 return None62 if base_model_id == self.base_model_id and task_name == self.task_name and hasattr(63 self, 'pipe'):64 return self.pipe65 model_id = CONTROLNET_MODEL_IDS[task_name]66 controlnet = ControlNetModel.from_pretrained(model_id,67 torch_dtype=torch.float16)68 pipe = StableDiffusionControlNetPipeline.from_pretrained(69 base_model_id,70 safety_checker=None,71 controlnet=controlnet,72 torch_dtype=torch.float16)73 pipe.scheduler = UniPCMultistepScheduler.from_config(74 pipe.scheduler.config)75 pipe.enable_xformers_memory_efficient_attention()76 pipe.to(self.device)77 torch.cuda.empty_cache()78 gc.collect()79 self.base_model_id = base_model_id80 self.task_name = task_name81 return pipe82 83 def set_base_model(self, base_model_id: str) -> str:84 if not base_model_id or base_model_id == self.base_model_id:85 return self.base_model_id86 del self.pipe87 torch.cuda.empty_cache()88 gc.collect()89 try:90 self.pipe = self.load_pipe(base_model_id, self.task_name)91 except Exception:92 self.pipe = self.load_pipe(self.base_model_id, self.task_name)93 return self.base_model_id94 95 def load_controlnet_weight(self, task_name: str) -> None:96 if task_name == self.task_name:97 return98 if 'controlnet' in self.pipe.__dict__:99 del self.pipe.controlnet100 torch.cuda.empty_cache()101 gc.collect()102 model_id = CONTROLNET_MODEL_IDS[task_name]103 controlnet = ControlNetModel.from_pretrained(model_id,104 torch_dtype=torch.float16)105 controlnet.to(self.device)106 torch.cuda.empty_cache()107 gc.collect()108 self.pipe.controlnet = controlnet109 self.task_name = task_name110 111 def get_prompt(self, prompt: str, additional_prompt: str) -> str:112 if not prompt:113 prompt = additional_prompt114 else:115 prompt = f'{prompt}, {additional_prompt}'116 return prompt117 118 @torch.autocast('cuda')119 def run_pipe(120 self,121 prompt: str,122 negative_prompt: str,123 control_image: PIL.Image.Image,124 num_images: int,125 num_steps: int,126 guidance_scale: float,127 seed: int,128 ) -> list[PIL.Image.Image]:129 if seed == -1:130 seed = np.random.randint(0, np.iinfo(np.int64).max)131 generator = torch.Generator().manual_seed(seed)132 return self.pipe(prompt=prompt,133 negative_prompt=negative_prompt,134 guidance_scale=guidance_scale,135 num_images_per_prompt=num_images,136 num_inference_steps=num_steps,137 generator=generator,138 image=control_image).images139 140 @staticmethod141 def preprocess_canny(142 input_image: np.ndarray,143 image_resolution: int,144 low_threshold: int,145 high_threshold: int,146 ) -> tuple[PIL.Image.Image, PIL.Image.Image]:147 image = resize_image(HWC3(input_image), image_resolution)148 control_image = apply_canny(image, low_threshold, high_threshold)149 control_image = HWC3(control_image)150 vis_control_image = 255 - control_image151 return PIL.Image.fromarray(control_image), PIL.Image.fromarray(152 vis_control_image)153 154 @torch.inference_mode()155 def process_canny(156 self,157 input_image: np.ndarray,158 prompt: str,159 additional_prompt: str,160 negative_prompt: str,161 num_images: int,162 image_resolution: int,163 num_steps: int,164 guidance_scale: float,165 seed: int,166 low_threshold: int,167 high_threshold: int,168 ) -> list[PIL.Image.Image]:169 control_image, vis_control_image = self.preprocess_canny(170 input_image=input_image,171 image_resolution=image_resolution,172 low_threshold=low_threshold,173 high_threshold=high_threshold,174 )175 self.load_controlnet_weight('canny')176 results = self.run_pipe(177 prompt=self.get_prompt(prompt, additional_prompt),178 negative_prompt=negative_prompt,179 control_image=control_image,180 num_images=num_images,181 num_steps=num_steps,182 guidance_scale=guidance_scale,183 seed=seed,184 )185 return [vis_control_image] + results186 187 @staticmethod188 def preprocess_hough(189 input_image: np.ndarray,190 image_resolution: int,191 detect_resolution: int,192 value_threshold: float,193 distance_threshold: float,194 ) -> tuple[PIL.Image.Image, PIL.Image.Image]:195 input_image = HWC3(input_image)196 control_image = apply_mlsd(197 resize_image(input_image, detect_resolution), value_threshold,198 distance_threshold)199 control_image = HWC3(control_image)200 image = resize_image(input_image, image_resolution)201 H, W = image.shape[:2]202 control_image = cv2.resize(control_image, (W, H),203 interpolation=cv2.INTER_NEAREST)204 205 vis_control_image = 255 - cv2.dilate(206 control_image, np.ones(shape=(3, 3), dtype=np.uint8), iterations=1)207 208 return PIL.Image.fromarray(control_image), PIL.Image.fromarray(209 vis_control_image)210 211 @torch.inference_mode()212 def process_hough(213 self,214 input_image: np.ndarray,215 prompt: str,216 additional_prompt: str,217 negative_prompt: str,218 num_images: int,219 image_resolution: int,220 detect_resolution: int,221 num_steps: int,222 guidance_scale: float,223 seed: int,224 value_threshold: float,225 distance_threshold: float,226 ) -> list[PIL.Image.Image]:227 control_image, vis_control_image = self.preprocess_hough(228 input_image=input_image,229 image_resolution=image_resolution,230 detect_resolution=detect_resolution,231 value_threshold=value_threshold,232 distance_threshold=distance_threshold,233 )234 self.load_controlnet_weight('hough')235 results = self.run_pipe(236 prompt=self.get_prompt(prompt, additional_prompt),237 negative_prompt=negative_prompt,238 control_image=control_image,239 num_images=num_images,240 num_steps=num_steps,241 guidance_scale=guidance_scale,242 seed=seed,243 )244 return [vis_control_image] + results245 246 @staticmethod247 def preprocess_hed(248 input_image: np.ndarray,249 image_resolution: int,250 detect_resolution: int,251 ) -> tuple[PIL.Image.Image, PIL.Image.Image]:252 input_image = HWC3(input_image)253 control_image = apply_hed(resize_image(input_image, detect_resolution))254 control_image = HWC3(control_image)255 image = resize_image(input_image, image_resolution)256 H, W = image.shape[:2]257 control_image = cv2.resize(control_image, (W, H),258 interpolation=cv2.INTER_LINEAR)259 return PIL.Image.fromarray(control_image), PIL.Image.fromarray(260 control_image)261 262 @torch.inference_mode()263 def process_hed(264 self,265 input_image: np.ndarray,266 prompt: str,267 additional_prompt: str,268 negative_prompt: str,269 num_images: int,270 image_resolution: int,271 detect_resolution: int,272 num_steps: int,273 guidance_scale: float,274 seed: int,275 ) -> list[PIL.Image.Image]:276 control_image, vis_control_image = self.preprocess_hed(277 input_image=input_image,278 image_resolution=image_resolution,279 detect_resolution=detect_resolution,280 )281 self.load_controlnet_weight('hed')282 results = self.run_pipe(283 prompt=self.get_prompt(prompt, additional_prompt),284 negative_prompt=negative_prompt,285 control_image=control_image,286 num_images=num_images,287 num_steps=num_steps,288 guidance_scale=guidance_scale,289 seed=seed,290 )291 return [vis_control_image] + results292 293 @staticmethod294 def preprocess_scribble(295 input_image: np.ndarray,296 image_resolution: int,297 ) -> tuple[PIL.Image.Image, PIL.Image.Image]:298 image = resize_image(HWC3(input_image), image_resolution)299 control_image = np.zeros_like(image, dtype=np.uint8)300 control_image[np.min(image, axis=2) < 127] = 255301 vis_control_image = 255 - control_image302 return PIL.Image.fromarray(control_image), PIL.Image.fromarray(303 vis_control_image)304 305 @torch.inference_mode()306 def process_scribble(307 self,308 input_image: np.ndarray,309 prompt: str,310 additional_prompt: str,311 negative_prompt: str,312 num_images: int,313 image_resolution: int,314 num_steps: int,315 guidance_scale: float,316 seed: int,317 ) -> list[PIL.Image.Image]:318 control_image, vis_control_image = self.preprocess_scribble(319 input_image=input_image,320 image_resolution=image_resolution,321 )322 self.load_controlnet_weight('scribble')323 results = self.run_pipe(324 prompt=self.get_prompt(prompt, additional_prompt),325 negative_prompt=negative_prompt,326 control_image=control_image,327 num_images=num_images,328 num_steps=num_steps,329 guidance_scale=guidance_scale,330 seed=seed,331 )332 return [vis_control_image] + results333 334 @staticmethod335 def preprocess_scribble_interactive(336 input_image: np.ndarray,337 image_resolution: int,338 ) -> tuple[PIL.Image.Image, PIL.Image.Image]:339 image = resize_image(HWC3(input_image['mask'][:, :, 0]),340 image_resolution)341 control_image = np.zeros_like(image, dtype=np.uint8)342 control_image[np.min(image, axis=2) > 127] = 255343 vis_control_image = 255 - control_image344 return PIL.Image.fromarray(control_image), PIL.Image.fromarray(345 vis_control_image)346 347 @torch.inference_mode()348 def process_scribble_interactive(349 self,350 input_image: np.ndarray,351 prompt: str,352 additional_prompt: str,353 negative_prompt: str,354 num_images: int,355 image_resolution: int,356 num_steps: int,357 guidance_scale: float,358 seed: int,359 ) -> list[PIL.Image.Image]:360 control_image, vis_control_image = self.preprocess_scribble_interactive(361 input_image=input_image,362 image_resolution=image_resolution,363 )364 self.load_controlnet_weight('scribble')365 results = self.run_pipe(366 prompt=self.get_prompt(prompt, additional_prompt),367 negative_prompt=negative_prompt,368 control_image=control_image,369 num_images=num_images,370 num_steps=num_steps,371 guidance_scale=guidance_scale,372 seed=seed,373 )374 return [vis_control_image] + results375 376 @staticmethod377 def preprocess_fake_scribble(378 input_image: np.ndarray,379 image_resolution: int,380 detect_resolution: int,381 ) -> tuple[PIL.Image.Image, PIL.Image.Image]:382 input_image = HWC3(input_image)383 control_image = apply_hed(resize_image(input_image, detect_resolution))384 control_image = HWC3(control_image)385 image = resize_image(input_image, image_resolution)386 H, W = image.shape[:2]387 388 control_image = cv2.resize(control_image, (W, H),389 interpolation=cv2.INTER_LINEAR)390 control_image = nms(control_image, 127, 3.0)391 control_image = cv2.GaussianBlur(control_image, (0, 0), 3.0)392 control_image[control_image > 4] = 255393 control_image[control_image < 255] = 0394 395 vis_control_image = 255 - control_image396 397 return PIL.Image.fromarray(control_image), PIL.Image.fromarray(398 vis_control_image)399 400 @torch.inference_mode()401 def process_fake_scribble(402 self,403 input_image: np.ndarray,404 prompt: str,405 additional_prompt: str,406 negative_prompt: str,407 num_images: int,408 image_resolution: int,409 detect_resolution: int,410 num_steps: int,411 guidance_scale: float,412 seed: int,413 ) -> list[PIL.Image.Image]:414 control_image, vis_control_image = self.preprocess_fake_scribble(415 input_image=input_image,416 image_resolution=image_resolution,417 detect_resolution=detect_resolution,418 )419 self.load_controlnet_weight('scribble')420 results = self.run_pipe(421 prompt=self.get_prompt(prompt, additional_prompt),422 negative_prompt=negative_prompt,423 control_image=control_image,424 num_images=num_images,425 num_steps=num_steps,426 guidance_scale=guidance_scale,427 seed=seed,428 )429 return [vis_control_image] + results430 431 @staticmethod432 def preprocess_pose(433 input_image: np.ndarray,434 image_resolution: int,435 detect_resolution: int,436 is_pose_image: bool,437 ) -> tuple[PIL.Image.Image, PIL.Image.Image]:438 input_image = HWC3(input_image)439 if not is_pose_image:440 control_image, _ = apply_openpose(441 resize_image(input_image, detect_resolution))442 control_image = HWC3(control_image)443 image = resize_image(input_image, image_resolution)444 H, W = image.shape[:2]445 control_image = cv2.resize(control_image, (W, H),446 interpolation=cv2.INTER_NEAREST)447 else:448 control_image = resize_image(input_image, image_resolution)449 450 return PIL.Image.fromarray(control_image), PIL.Image.fromarray(451 control_image)452 453 @torch.inference_mode()454 def process_pose(455 self,456 input_image: np.ndarray,457 prompt: str,458 additional_prompt: str,459 negative_prompt: str,460 num_images: int,461 image_resolution: int,462 detect_resolution: int,463 num_steps: int,464 guidance_scale: float,465 seed: int,466 is_pose_image: bool,467 ) -> list[PIL.Image.Image]:468 control_image, vis_control_image = self.preprocess_pose(469 input_image=input_image,470 image_resolution=image_resolution,471 detect_resolution=detect_resolution,472 is_pose_image=is_pose_image,473 )474 self.load_controlnet_weight('pose')475 results = self.run_pipe(476 prompt=self.get_prompt(prompt, additional_prompt),477 negative_prompt=negative_prompt,478 control_image=control_image,479 num_images=num_images,480 num_steps=num_steps,481 guidance_scale=guidance_scale,482 seed=seed,483 )484 return [vis_control_image] + results485 486 @staticmethod487 def preprocess_seg(488 input_image: np.ndarray,489 image_resolution: int,490 detect_resolution: int,491 is_segmentation_map: bool,492 ) -> tuple[PIL.Image.Image, PIL.Image.Image]:493 input_image = HWC3(input_image)494 if not is_segmentation_map:495 control_image = apply_uniformer(496 resize_image(input_image, detect_resolution))497 image = resize_image(input_image, image_resolution)498 H, W = image.shape[:2]499 control_image = cv2.resize(control_image, (W, H),500 interpolation=cv2.INTER_NEAREST)501 else:502 control_image = resize_image(input_image, image_resolution)503 return PIL.Image.fromarray(control_image), PIL.Image.fromarray(504 control_image)505 506 @torch.inference_mode()507 def process_seg(508 self,509 input_image: np.ndarray,510 prompt: str,511 additional_prompt: str,512 negative_prompt: str,513 num_images: int,514 image_resolution: int,515 detect_resolution: int,516 num_steps: int,517 guidance_scale: float,518 seed: int,519 is_segmentation_map: bool,520 ) -> list[PIL.Image.Image]:521 control_image, vis_control_image = self.preprocess_seg(522 input_image=input_image,523 image_resolution=image_resolution,524 detect_resolution=detect_resolution,525 is_segmentation_map=is_segmentation_map,526 )527 self.load_controlnet_weight('seg')528 results = self.run_pipe(529 prompt=self.get_prompt(prompt, additional_prompt),530 negative_prompt=negative_prompt,531 control_image=control_image,532 num_images=num_images,533 num_steps=num_steps,534 guidance_scale=guidance_scale,535 seed=seed,536 )537 return [vis_control_image] + results538 539 @staticmethod540 def preprocess_depth(541 input_image: np.ndarray,542 image_resolution: int,543 detect_resolution: int,544 is_depth_image: bool,545 ) -> tuple[PIL.Image.Image, PIL.Image.Image]:546 input_image = HWC3(input_image)547 if not is_depth_image:548 control_image, _ = apply_midas(549 resize_image(input_image, detect_resolution))550 control_image = HWC3(control_image)551 image = resize_image(input_image, image_resolution)552 H, W = image.shape[:2]553 control_image = cv2.resize(control_image, (W, H),554 interpolation=cv2.INTER_LINEAR)555 else:556 control_image = resize_image(input_image, image_resolution)557 return PIL.Image.fromarray(control_image), PIL.Image.fromarray(558 control_image)559 560 @torch.inference_mode()561 def process_depth(562 self,563 input_image: np.ndarray,564 prompt: str,565 additional_prompt: str,566 negative_prompt: str,567 num_images: int,568 image_resolution: int,569 detect_resolution: int,570 num_steps: int,571 guidance_scale: float,572 seed: int,573 is_depth_image: bool,574 ) -> list[PIL.Image.Image]:575 control_image, vis_control_image = self.preprocess_depth(576 input_image=input_image,577 image_resolution=image_resolution,578 detect_resolution=detect_resolution,579 is_depth_image=is_depth_image,580 )581 self.load_controlnet_weight('depth')582 results = self.run_pipe(583 prompt=self.get_prompt(prompt, additional_prompt),584 negative_prompt=negative_prompt,585 control_image=control_image,586 num_images=num_images,587 num_steps=num_steps,588 guidance_scale=guidance_scale,589 seed=seed,590 )591 return [vis_control_image] + results592 593 @staticmethod594 def preprocess_normal(595 input_image: np.ndarray,596 image_resolution: int,597 detect_resolution: int,598 bg_threshold: float,599 is_normal_image: bool,600 ) -> tuple[PIL.Image.Image, PIL.Image.Image]:601 input_image = HWC3(input_image)602 if not is_normal_image:603 _, control_image = apply_midas(resize_image(604 input_image, detect_resolution),605 bg_th=bg_threshold)606 control_image = HWC3(control_image)607 image = resize_image(input_image, image_resolution)608 H, W = image.shape[:2]609 control_image = cv2.resize(control_image, (W, H),610 interpolation=cv2.INTER_LINEAR)611 else:612 control_image = resize_image(input_image, image_resolution)613 return PIL.Image.fromarray(control_image), PIL.Image.fromarray(614 control_image)615 616 @torch.inference_mode()617 def process_normal(618 self,619 input_image: np.ndarray,620 prompt: str,621 additional_prompt: str,622 negative_prompt: str,623 num_images: int,624 image_resolution: int,625 detect_resolution: int,626 num_steps: int,627 guidance_scale: float,628 seed: int,629 bg_threshold: float,630 is_normal_image: bool,631 ) -> list[PIL.Image.Image]:632 control_image, vis_control_image = self.preprocess_normal(633 input_image=input_image,634 image_resolution=image_resolution,635 detect_resolution=detect_resolution,636 bg_threshold=bg_threshold,637 is_normal_image=is_normal_image,638 )639 self.load_controlnet_weight('normal')640 results = self.run_pipe(641 prompt=self.get_prompt(prompt, additional_prompt),642 negative_prompt=negative_prompt,643 control_image=control_image,644 num_images=num_images,645 num_steps=num_steps,646 guidance_scale=guidance_scale,647 seed=seed,648 )649 return [vis_control_image] + results650