HIST/ControlNet
0
1# This file is adapted from gradio_*.py in https://github.com/lllyasviel/ControlNet/tree/f4748e3630d8141d7765e2bd9b1e348f478477072# The original license file is LICENSE.ControlNet this repo.3from __future__ import annotations4 5import pathlib6import random7import shlex8import subprocess9import sys10 11import cv212import einops13import numpy as np14import torch15from pytorch_lightning import seed_everything16 17sys.path.append('ControlNet')18 19import config20from annotator.canny import apply_canny21from annotator.hed import apply_hed, nms22from annotator.midas import apply_midas23from annotator.mlsd import apply_mlsd24from annotator.openpose import apply_openpose25from annotator.uniformer import apply_uniformer26from annotator.util import HWC3, resize_image27from cldm.model import create_model, load_state_dict28from ldm.models.diffusion.ddim import DDIMSampler29from share import *30 31 32class Model:33 WEIGHT_NAMES = {34 'canny': 'control_sd15_canny.pth',35 'hough': 'control_sd15_mlsd.pth',36 'hed': 'control_sd15_hed.pth',37 'scribble': 'control_sd15_scribble.pth',38 'pose': 'control_sd15_openpose.pth',39 'seg': 'control_sd15_seg.pth',40 'depth': 'control_sd15_depth.pth',41 'normal': 'control_sd15_normal.pth',42 }43 44 def __init__(self,45 model_config_path: str = 'ControlNet/models/cldm_v15.yaml',46 model_dir: str = 'models'):47 self.device = torch.device(48 'cuda:0' if torch.cuda.is_available() else 'cpu')49 self.model = create_model(model_config_path).to(self.device)50 self.ddim_sampler = DDIMSampler(self.model)51 self.task_name = ''52 53 self.model_dir = pathlib.Path(model_dir)54 self.download_models()55 56 def load_weight(self, task_name: str) -> None:57 if task_name == self.task_name:58 return59 weight_path = self.get_weight_path(task_name)60 self.model.load_state_dict(61 load_state_dict(weight_path, location=self.device))62 self.task_name = task_name63 64 def get_weight_path(self, task_name: str) -> str:65 if 'scribble' in task_name:66 task_name = 'scribble'67 return f'{self.model_dir}/{self.WEIGHT_NAMES[task_name]}'68 69 def download_models(self):70 self.model_dir.mkdir(exist_ok=True, parents=True)71 for name in self.WEIGHT_NAMES.values():72 out_path = self.model_dir / name73 if out_path.exists():74 continue75 subprocess.run(76 shlex.split(77 f'wget https://huggingface.co/ckpt/ControlNet/resolve/main/{name} -O {out_path}'78 ))79 80 @torch.inference_mode()81 def process_canny(self, input_image, prompt, a_prompt, n_prompt,82 num_samples, image_resolution, ddim_steps, scale, seed,83 eta, low_threshold, high_threshold):84 self.load_weight('canny')85 86 img = resize_image(HWC3(input_image), image_resolution)87 H, W, C = img.shape88 89 detected_map = apply_canny(img, low_threshold, high_threshold)90 detected_map = HWC3(detected_map)91 92 control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.093 control = torch.stack([control for _ in range(num_samples)], dim=0)94 control = einops.rearrange(control, 'b h w c -> b c h w').clone()95 96 if seed == -1:97 seed = random.randint(0, 65535)98 seed_everything(seed)99 100 if config.save_memory:101 self.model.low_vram_shift(is_diffusing=False)102 103 cond = {104 'c_concat': [control],105 'c_crossattn': [106 self.model.get_learned_conditioning(107 [prompt + ', ' + a_prompt] * num_samples)108 ]109 }110 un_cond = {111 'c_concat': [control],112 'c_crossattn':113 [self.model.get_learned_conditioning([n_prompt] * num_samples)]114 }115 shape = (4, H // 8, W // 8)116 117 if config.save_memory:118 self.model.low_vram_shift(is_diffusing=True)119 120 samples, intermediates = self.ddim_sampler.sample(121 ddim_steps,122 num_samples,123 shape,124 cond,125 verbose=False,126 eta=eta,127 unconditional_guidance_scale=scale,128 unconditional_conditioning=un_cond)129 130 if config.save_memory:131 self.model.low_vram_shift(is_diffusing=False)132 133 x_samples = self.model.decode_first_stage(samples)134 x_samples = (135 einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +136 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)137 138 results = [x_samples[i] for i in range(num_samples)]139 return [255 - detected_map] + results140 141 @torch.inference_mode()142 def process_hough(self, input_image, prompt, a_prompt, n_prompt,143 num_samples, image_resolution, detect_resolution,144 ddim_steps, scale, seed, eta, value_threshold,145 distance_threshold):146 self.load_weight('hough')147 148 input_image = HWC3(input_image)149 detected_map = apply_mlsd(resize_image(input_image, detect_resolution),150 value_threshold, distance_threshold)151 detected_map = HWC3(detected_map)152 img = resize_image(input_image, image_resolution)153 H, W, C = img.shape154 155 detected_map = cv2.resize(detected_map, (W, H),156 interpolation=cv2.INTER_NEAREST)157 158 control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0159 control = torch.stack([control for _ in range(num_samples)], dim=0)160 control = einops.rearrange(control, 'b h w c -> b c h w').clone()161 162 if seed == -1:163 seed = random.randint(0, 65535)164 seed_everything(seed)165 166 if config.save_memory:167 self.model.low_vram_shift(is_diffusing=False)168 169 cond = {170 'c_concat': [control],171 'c_crossattn': [172 self.model.get_learned_conditioning(173 [prompt + ', ' + a_prompt] * num_samples)174 ]175 }176 un_cond = {177 'c_concat': [control],178 'c_crossattn':179 [self.model.get_learned_conditioning([n_prompt] * num_samples)]180 }181 shape = (4, H // 8, W // 8)182 183 if config.save_memory:184 self.model.low_vram_shift(is_diffusing=True)185 186 samples, intermediates = self.ddim_sampler.sample(187 ddim_steps,188 num_samples,189 shape,190 cond,191 verbose=False,192 eta=eta,193 unconditional_guidance_scale=scale,194 unconditional_conditioning=un_cond)195 196 if config.save_memory:197 self.model.low_vram_shift(is_diffusing=False)198 199 x_samples = self.model.decode_first_stage(samples)200 x_samples = (201 einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +202 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)203 204 results = [x_samples[i] for i in range(num_samples)]205 return [206 255 - cv2.dilate(detected_map,207 np.ones(shape=(3, 3), dtype=np.uint8),208 iterations=1)209 ] + results210 211 @torch.inference_mode()212 def process_hed(self, input_image, prompt, a_prompt, n_prompt, num_samples,213 image_resolution, detect_resolution, ddim_steps, scale,214 seed, eta):215 self.load_weight('hed')216 217 input_image = HWC3(input_image)218 detected_map = apply_hed(resize_image(input_image, detect_resolution))219 detected_map = HWC3(detected_map)220 img = resize_image(input_image, image_resolution)221 H, W, C = img.shape222 223 detected_map = cv2.resize(detected_map, (W, H),224 interpolation=cv2.INTER_LINEAR)225 226 control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0227 control = torch.stack([control for _ in range(num_samples)], dim=0)228 control = einops.rearrange(control, 'b h w c -> b c h w').clone()229 230 if seed == -1:231 seed = random.randint(0, 65535)232 seed_everything(seed)233 234 if config.save_memory:235 self.model.low_vram_shift(is_diffusing=False)236 237 cond = {238 'c_concat': [control],239 'c_crossattn': [240 self.model.get_learned_conditioning(241 [prompt + ', ' + a_prompt] * num_samples)242 ]243 }244 un_cond = {245 'c_concat': [control],246 'c_crossattn':247 [self.model.get_learned_conditioning([n_prompt] * num_samples)]248 }249 shape = (4, H // 8, W // 8)250 251 if config.save_memory:252 self.model.low_vram_shift(is_diffusing=True)253 254 samples, intermediates = self.ddim_sampler.sample(255 ddim_steps,256 num_samples,257 shape,258 cond,259 verbose=False,260 eta=eta,261 unconditional_guidance_scale=scale,262 unconditional_conditioning=un_cond)263 264 if config.save_memory:265 self.model.low_vram_shift(is_diffusing=False)266 267 x_samples = self.model.decode_first_stage(samples)268 x_samples = (269 einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +270 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)271 272 results = [x_samples[i] for i in range(num_samples)]273 return [detected_map] + results274 275 @torch.inference_mode()276 def process_scribble(self, input_image, prompt, a_prompt, n_prompt,277 num_samples, image_resolution, ddim_steps, scale,278 seed, eta):279 self.load_weight('scribble')280 281 img = resize_image(HWC3(input_image), image_resolution)282 H, W, C = img.shape283 284 detected_map = np.zeros_like(img, dtype=np.uint8)285 detected_map[np.min(img, axis=2) < 127] = 255286 287 control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0288 control = torch.stack([control for _ in range(num_samples)], dim=0)289 control = einops.rearrange(control, 'b h w c -> b c h w').clone()290 291 if seed == -1:292 seed = random.randint(0, 65535)293 seed_everything(seed)294 295 if config.save_memory:296 self.model.low_vram_shift(is_diffusing=False)297 298 cond = {299 'c_concat': [control],300 'c_crossattn': [301 self.model.get_learned_conditioning(302 [prompt + ', ' + a_prompt] * num_samples)303 ]304 }305 un_cond = {306 'c_concat': [control],307 'c_crossattn':308 [self.model.get_learned_conditioning([n_prompt] * num_samples)]309 }310 shape = (4, H // 8, W // 8)311 312 if config.save_memory:313 self.model.low_vram_shift(is_diffusing=True)314 315 samples, intermediates = self.ddim_sampler.sample(316 ddim_steps,317 num_samples,318 shape,319 cond,320 verbose=False,321 eta=eta,322 unconditional_guidance_scale=scale,323 unconditional_conditioning=un_cond)324 325 if config.save_memory:326 self.model.low_vram_shift(is_diffusing=False)327 328 x_samples = self.model.decode_first_stage(samples)329 x_samples = (330 einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +331 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)332 333 results = [x_samples[i] for i in range(num_samples)]334 return [255 - detected_map] + results335 336 @torch.inference_mode()337 def process_scribble_interactive(self, input_image, prompt, a_prompt,338 n_prompt, num_samples, image_resolution,339 ddim_steps, scale, seed, eta):340 self.load_weight('scribble')341 342 img = resize_image(HWC3(input_image['mask'][:, :, 0]),343 image_resolution)344 H, W, C = img.shape345 346 detected_map = np.zeros_like(img, dtype=np.uint8)347 detected_map[np.min(img, axis=2) > 127] = 255348 349 control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0350 control = torch.stack([control for _ in range(num_samples)], dim=0)351 control = einops.rearrange(control, 'b h w c -> b c h w').clone()352 353 if seed == -1:354 seed = random.randint(0, 65535)355 seed_everything(seed)356 357 if config.save_memory:358 self.model.low_vram_shift(is_diffusing=False)359 360 cond = {361 'c_concat': [control],362 'c_crossattn': [363 self.model.get_learned_conditioning(364 [prompt + ', ' + a_prompt] * num_samples)365 ]366 }367 un_cond = {368 'c_concat': [control],369 'c_crossattn':370 [self.model.get_learned_conditioning([n_prompt] * num_samples)]371 }372 shape = (4, H // 8, W // 8)373 374 if config.save_memory:375 self.model.low_vram_shift(is_diffusing=True)376 377 samples, intermediates = self.ddim_sampler.sample(378 ddim_steps,379 num_samples,380 shape,381 cond,382 verbose=False,383 eta=eta,384 unconditional_guidance_scale=scale,385 unconditional_conditioning=un_cond)386 387 if config.save_memory:388 self.model.low_vram_shift(is_diffusing=False)389 390 x_samples = self.model.decode_first_stage(samples)391 x_samples = (392 einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +393 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)394 395 results = [x_samples[i] for i in range(num_samples)]396 return [255 - detected_map] + results397 398 @torch.inference_mode()399 def process_fake_scribble(self, input_image, prompt, a_prompt, n_prompt,400 num_samples, image_resolution, detect_resolution,401 ddim_steps, scale, seed, eta):402 self.load_weight('scribble')403 404 input_image = HWC3(input_image)405 detected_map = apply_hed(resize_image(input_image, detect_resolution))406 detected_map = HWC3(detected_map)407 img = resize_image(input_image, image_resolution)408 H, W, C = img.shape409 410 detected_map = cv2.resize(detected_map, (W, H),411 interpolation=cv2.INTER_LINEAR)412 detected_map = nms(detected_map, 127, 3.0)413 detected_map = cv2.GaussianBlur(detected_map, (0, 0), 3.0)414 detected_map[detected_map > 4] = 255415 detected_map[detected_map < 255] = 0416 417 control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0418 control = torch.stack([control for _ in range(num_samples)], dim=0)419 control = einops.rearrange(control, 'b h w c -> b c h w').clone()420 421 if seed == -1:422 seed = random.randint(0, 65535)423 seed_everything(seed)424 425 if config.save_memory:426 self.model.low_vram_shift(is_diffusing=False)427 428 cond = {429 'c_concat': [control],430 'c_crossattn': [431 self.model.get_learned_conditioning(432 [prompt + ', ' + a_prompt] * num_samples)433 ]434 }435 un_cond = {436 'c_concat': [control],437 'c_crossattn':438 [self.model.get_learned_conditioning([n_prompt] * num_samples)]439 }440 shape = (4, H // 8, W // 8)441 442 if config.save_memory:443 self.model.low_vram_shift(is_diffusing=True)444 445 samples, intermediates = self.ddim_sampler.sample(446 ddim_steps,447 num_samples,448 shape,449 cond,450 verbose=False,451 eta=eta,452 unconditional_guidance_scale=scale,453 unconditional_conditioning=un_cond)454 455 if config.save_memory:456 self.model.low_vram_shift(is_diffusing=False)457 458 x_samples = self.model.decode_first_stage(samples)459 x_samples = (460 einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +461 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)462 463 results = [x_samples[i] for i in range(num_samples)]464 return [255 - detected_map] + results465 466 @torch.inference_mode()467 def process_pose(self, input_image, prompt, a_prompt, n_prompt,468 num_samples, image_resolution, detect_resolution,469 ddim_steps, scale, seed, eta):470 self.load_weight('pose')471 472 input_image = HWC3(input_image)473 detected_map, _ = apply_openpose(474 resize_image(input_image, detect_resolution))475 detected_map = HWC3(detected_map)476 img = resize_image(input_image, image_resolution)477 H, W, C = img.shape478 479 detected_map = cv2.resize(detected_map, (W, H),480 interpolation=cv2.INTER_NEAREST)481 482 control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0483 control = torch.stack([control for _ in range(num_samples)], dim=0)484 control = einops.rearrange(control, 'b h w c -> b c h w').clone()485 486 if seed == -1:487 seed = random.randint(0, 65535)488 seed_everything(seed)489 490 if config.save_memory:491 self.model.low_vram_shift(is_diffusing=False)492 493 cond = {494 'c_concat': [control],495 'c_crossattn': [496 self.model.get_learned_conditioning(497 [prompt + ', ' + a_prompt] * num_samples)498 ]499 }500 un_cond = {501 'c_concat': [control],502 'c_crossattn':503 [self.model.get_learned_conditioning([n_prompt] * num_samples)]504 }505 shape = (4, H // 8, W // 8)506 507 if config.save_memory:508 self.model.low_vram_shift(is_diffusing=True)509 510 samples, intermediates = self.ddim_sampler.sample(511 ddim_steps,512 num_samples,513 shape,514 cond,515 verbose=False,516 eta=eta,517 unconditional_guidance_scale=scale,518 unconditional_conditioning=un_cond)519 520 if config.save_memory:521 self.model.low_vram_shift(is_diffusing=False)522 523 x_samples = self.model.decode_first_stage(samples)524 x_samples = (525 einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +526 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)527 528 results = [x_samples[i] for i in range(num_samples)]529 return [detected_map] + results530 531 @torch.inference_mode()532 def process_seg(self, input_image, prompt, a_prompt, n_prompt, num_samples,533 image_resolution, detect_resolution, ddim_steps, scale,534 seed, eta):535 self.load_weight('seg')536 537 input_image = HWC3(input_image)538 detected_map = apply_uniformer(539 resize_image(input_image, detect_resolution))540 img = resize_image(input_image, image_resolution)541 H, W, C = img.shape542 543 detected_map = cv2.resize(detected_map, (W, H),544 interpolation=cv2.INTER_NEAREST)545 546 control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0547 control = torch.stack([control for _ in range(num_samples)], dim=0)548 control = einops.rearrange(control, 'b h w c -> b c h w').clone()549 550 if seed == -1:551 seed = random.randint(0, 65535)552 seed_everything(seed)553 554 if config.save_memory:555 self.model.low_vram_shift(is_diffusing=False)556 557 cond = {558 'c_concat': [control],559 'c_crossattn': [560 self.model.get_learned_conditioning(561 [prompt + ', ' + a_prompt] * num_samples)562 ]563 }564 un_cond = {565 'c_concat': [control],566 'c_crossattn':567 [self.model.get_learned_conditioning([n_prompt] * num_samples)]568 }569 shape = (4, H // 8, W // 8)570 571 if config.save_memory:572 self.model.low_vram_shift(is_diffusing=True)573 574 samples, intermediates = self.ddim_sampler.sample(575 ddim_steps,576 num_samples,577 shape,578 cond,579 verbose=False,580 eta=eta,581 unconditional_guidance_scale=scale,582 unconditional_conditioning=un_cond)583 584 if config.save_memory:585 self.model.low_vram_shift(is_diffusing=False)586 587 x_samples = self.model.decode_first_stage(samples)588 x_samples = (589 einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +590 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)591 592 results = [x_samples[i] for i in range(num_samples)]593 return [detected_map] + results594 595 @torch.inference_mode()596 def process_depth(self, input_image, prompt, a_prompt, n_prompt,597 num_samples, image_resolution, detect_resolution,598 ddim_steps, scale, seed, eta):599 self.load_weight('depth')600 601 input_image = HWC3(input_image)602 detected_map, _ = apply_midas(603 resize_image(input_image, detect_resolution))604 detected_map = HWC3(detected_map)605 img = resize_image(input_image, image_resolution)606 H, W, C = img.shape607 608 detected_map = cv2.resize(detected_map, (W, H),609 interpolation=cv2.INTER_LINEAR)610 611 control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0612 control = torch.stack([control for _ in range(num_samples)], dim=0)613 control = einops.rearrange(control, 'b h w c -> b c h w').clone()614 615 if seed == -1:616 seed = random.randint(0, 65535)617 seed_everything(seed)618 619 if config.save_memory:620 self.model.low_vram_shift(is_diffusing=False)621 622 cond = {623 'c_concat': [control],624 'c_crossattn': [625 self.model.get_learned_conditioning(626 [prompt + ', ' + a_prompt] * num_samples)627 ]628 }629 un_cond = {630 'c_concat': [control],631 'c_crossattn':632 [self.model.get_learned_conditioning([n_prompt] * num_samples)]633 }634 shape = (4, H // 8, W // 8)635 636 if config.save_memory:637 self.model.low_vram_shift(is_diffusing=True)638 639 samples, intermediates = self.ddim_sampler.sample(640 ddim_steps,641 num_samples,642 shape,643 cond,644 verbose=False,645 eta=eta,646 unconditional_guidance_scale=scale,647 unconditional_conditioning=un_cond)648 649 if config.save_memory:650 self.model.low_vram_shift(is_diffusing=False)651 652 x_samples = self.model.decode_first_stage(samples)653 x_samples = (654 einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +655 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)656 657 results = [x_samples[i] for i in range(num_samples)]658 return [detected_map] + results659 660 @torch.inference_mode()661 def process_normal(self, input_image, prompt, a_prompt, n_prompt,662 num_samples, image_resolution, detect_resolution,663 ddim_steps, scale, seed, eta, bg_threshold):664 self.load_weight('normal')665 666 input_image = HWC3(input_image)667 _, detected_map = apply_midas(resize_image(input_image,668 detect_resolution),669 bg_th=bg_threshold)670 detected_map = HWC3(detected_map)671 img = resize_image(input_image, image_resolution)672 H, W, C = img.shape673 674 detected_map = cv2.resize(detected_map, (W, H),675 interpolation=cv2.INTER_LINEAR)676 677 control = torch.from_numpy(678 detected_map[:, :, ::-1].copy()).float().cuda() / 255.0679 control = torch.stack([control for _ in range(num_samples)], dim=0)680 control = einops.rearrange(control, 'b h w c -> b c h w').clone()681 682 if seed == -1:683 seed = random.randint(0, 65535)684 seed_everything(seed)685 686 if config.save_memory:687 self.model.low_vram_shift(is_diffusing=False)688 689 cond = {690 'c_concat': [control],691 'c_crossattn': [692 self.model.get_learned_conditioning(693 [prompt + ', ' + a_prompt] * num_samples)694 ]695 }696 un_cond = {697 'c_concat': [control],698 'c_crossattn':699 [self.model.get_learned_conditioning([n_prompt] * num_samples)]700 }701 shape = (4, H // 8, W // 8)702 703 if config.save_memory:704 self.model.low_vram_shift(is_diffusing=True)705 706 samples, intermediates = self.ddim_sampler.sample(707 ddim_steps,708 num_samples,709 shape,710 cond,711 verbose=False,712 eta=eta,713 unconditional_guidance_scale=scale,714 unconditional_conditioning=un_cond)715 716 if config.save_memory:717 self.model.low_vram_shift(is_diffusing=False)718 719 x_samples = self.model.decode_first_stage(samples)720 x_samples = (721 einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +722 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)723 724 results = [x_samples[i] for i in range(num_samples)]725 return [detected_map] + results726