nivere/ControlNet-Video
1
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 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 31ORIGINAL_MODEL_NAMES = {32 'canny': 'control_sd15_canny.pth',33 'hough': 'control_sd15_mlsd.pth',34 'hed': 'control_sd15_hed.pth',35 'scribble': 'control_sd15_scribble.pth',36 'pose': 'control_sd15_openpose.pth',37 'seg': 'control_sd15_seg.pth',38 'depth': 'control_sd15_depth.pth',39 'normal': 'control_sd15_normal.pth',40}41ORIGINAL_WEIGHT_ROOT = 'https://huggingface.co/lllyasviel/ControlNet/resolve/main/models/'42 43LIGHTWEIGHT_MODEL_NAMES = {44 'canny': 'control_canny-fp16.safetensors',45 'hough': 'control_mlsd-fp16.safetensors',46 'hed': 'control_hed-fp16.safetensors',47 'scribble': 'control_scribble-fp16.safetensors',48 'pose': 'control_openpose-fp16.safetensors',49 'seg': 'control_seg-fp16.safetensors',50 'depth': 'control_depth-fp16.safetensors',51 'normal': 'control_normal-fp16.safetensors',52}53LIGHTWEIGHT_WEIGHT_ROOT = 'https://huggingface.co/webui/ControlNet-modules-safetensors/resolve/main/'54 55 56class Model:57 def __init__(self,58 model_config_path: str = 'ControlNet/models/cldm_v15.yaml',59 model_dir: str = 'models',60 use_lightweight: bool = True):61 self.device = torch.device(62 'cuda:0' if torch.cuda.is_available() else 'cpu')63 self.model = create_model(model_config_path).to(self.device)64 self.ddim_sampler = DDIMSampler(self.model)65 self.task_name = ''66 67 self.model_dir = pathlib.Path(model_dir)68 self.model_dir.mkdir(exist_ok=True, parents=True)69 70 self.use_lightweight = use_lightweight71 if use_lightweight:72 self.model_names = LIGHTWEIGHT_MODEL_NAMES73 self.weight_root = LIGHTWEIGHT_WEIGHT_ROOT74 base_model_url = 'https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors'75 self.load_base_model(base_model_url)76 else:77 self.model_names = ORIGINAL_MODEL_NAMES78 self.weight_root = ORIGINAL_WEIGHT_ROOT79 80 self.download_models()81 82 def download_base_model(self, model_url: str) -> pathlib.Path:83 model_name = model_url.split('/')[-1]84 out_path = self.model_dir / model_name85 if not out_path.exists():86 subprocess.run(shlex.split(f'wget {model_url} -O {out_path}'))87 return out_path88 89 def load_base_model(self, model_url: str) -> None:90 model_path = self.download_base_model(model_url)91 self.model.load_state_dict(load_state_dict(model_path,92 location=self.device.type),93 strict=False)94 95 def load_weight(self, task_name: str) -> None:96 if task_name == self.task_name:97 return98 weight_path = self.get_weight_path(task_name)99 if not self.use_lightweight:100 self.model.load_state_dict(101 load_state_dict(weight_path, location=self.device))102 else:103 self.model.control_model.load_state_dict(104 load_state_dict(weight_path, location=self.device.type))105 self.task_name = task_name106 107 def get_weight_path(self, task_name: str) -> str:108 if 'scribble' in task_name:109 task_name = 'scribble'110 return f'{self.model_dir}/{self.model_names[task_name]}'111 112 def download_models(self) -> None:113 self.model_dir.mkdir(exist_ok=True, parents=True)114 for name in self.model_names.values():115 out_path = self.model_dir / name116 if out_path.exists():117 continue118 subprocess.run(119 shlex.split(f'wget {self.weight_root}{name} -O {out_path}'))120 121 @torch.inference_mode()122 def process_canny(self, input_image, prompt, a_prompt, n_prompt,123 num_samples, image_resolution, ddim_steps, scale, seed,124 eta, low_threshold, high_threshold):125 self.load_weight('canny')126 127 img = resize_image(HWC3(input_image), image_resolution)128 H, W, C = img.shape129 130 detected_map = apply_canny(img, low_threshold, high_threshold)131 detected_map = HWC3(detected_map)132 133 control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0134 control = torch.stack([control for _ in range(num_samples)], dim=0)135 control = einops.rearrange(control, 'b h w c -> b c h w').clone()136 137 if seed == -1:138 seed = random.randint(0, 65535)139 seed_everything(seed)140 141 if config.save_memory:142 self.model.low_vram_shift(is_diffusing=False)143 144 cond = {145 'c_concat': [control],146 'c_crossattn': [147 self.model.get_learned_conditioning(148 [prompt + ', ' + a_prompt] * num_samples)149 ]150 }151 un_cond = {152 'c_concat': [control],153 'c_crossattn':154 [self.model.get_learned_conditioning([n_prompt] * num_samples)]155 }156 shape = (4, H // 8, W // 8)157 158 if config.save_memory:159 self.model.low_vram_shift(is_diffusing=True)160 161 samples, intermediates = self.ddim_sampler.sample(162 ddim_steps,163 num_samples,164 shape,165 cond,166 verbose=False,167 eta=eta,168 unconditional_guidance_scale=scale,169 unconditional_conditioning=un_cond)170 171 if config.save_memory:172 self.model.low_vram_shift(is_diffusing=False)173 174 x_samples = self.model.decode_first_stage(samples)175 x_samples = (176 einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +177 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)178 179 results = [x_samples[i] for i in range(num_samples)]180 return [255 - detected_map] + results181 182 @torch.inference_mode()183 def process_hough(self, input_image, prompt, a_prompt, n_prompt,184 num_samples, image_resolution, detect_resolution,185 ddim_steps, scale, seed, eta, value_threshold,186 distance_threshold):187 self.load_weight('hough')188 189 input_image = HWC3(input_image)190 detected_map = apply_mlsd(resize_image(input_image, detect_resolution),191 value_threshold, distance_threshold)192 detected_map = HWC3(detected_map)193 img = resize_image(input_image, image_resolution)194 H, W, C = img.shape195 196 detected_map = cv2.resize(detected_map, (W, H),197 interpolation=cv2.INTER_NEAREST)198 199 control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0200 control = torch.stack([control for _ in range(num_samples)], dim=0)201 control = einops.rearrange(control, 'b h w c -> b c h w').clone()202 203 if seed == -1:204 seed = random.randint(0, 65535)205 seed_everything(seed)206 207 if config.save_memory:208 self.model.low_vram_shift(is_diffusing=False)209 210 cond = {211 'c_concat': [control],212 'c_crossattn': [213 self.model.get_learned_conditioning(214 [prompt + ', ' + a_prompt] * num_samples)215 ]216 }217 un_cond = {218 'c_concat': [control],219 'c_crossattn':220 [self.model.get_learned_conditioning([n_prompt] * num_samples)]221 }222 shape = (4, H // 8, W // 8)223 224 if config.save_memory:225 self.model.low_vram_shift(is_diffusing=True)226 227 samples, intermediates = self.ddim_sampler.sample(228 ddim_steps,229 num_samples,230 shape,231 cond,232 verbose=False,233 eta=eta,234 unconditional_guidance_scale=scale,235 unconditional_conditioning=un_cond)236 237 if config.save_memory:238 self.model.low_vram_shift(is_diffusing=False)239 240 x_samples = self.model.decode_first_stage(samples)241 x_samples = (242 einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +243 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)244 245 results = [x_samples[i] for i in range(num_samples)]246 return [247 255 - cv2.dilate(detected_map,248 np.ones(shape=(3, 3), dtype=np.uint8),249 iterations=1)250 ] + results251 252 @torch.inference_mode()253 def process_hed(self, input_image, prompt, a_prompt, n_prompt, num_samples,254 image_resolution, detect_resolution, ddim_steps, scale,255 seed, eta):256 self.load_weight('hed')257 258 input_image = HWC3(input_image)259 detected_map = apply_hed(resize_image(input_image, detect_resolution))260 detected_map = HWC3(detected_map)261 img = resize_image(input_image, image_resolution)262 H, W, C = img.shape263 264 detected_map = cv2.resize(detected_map, (W, H),265 interpolation=cv2.INTER_LINEAR)266 267 control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0268 control = torch.stack([control for _ in range(num_samples)], dim=0)269 control = einops.rearrange(control, 'b h w c -> b c h w').clone()270 271 if seed == -1:272 seed = random.randint(0, 65535)273 seed_everything(seed)274 275 if config.save_memory:276 self.model.low_vram_shift(is_diffusing=False)277 278 cond = {279 'c_concat': [control],280 'c_crossattn': [281 self.model.get_learned_conditioning(282 [prompt + ', ' + a_prompt] * num_samples)283 ]284 }285 un_cond = {286 'c_concat': [control],287 'c_crossattn':288 [self.model.get_learned_conditioning([n_prompt] * num_samples)]289 }290 shape = (4, H // 8, W // 8)291 292 if config.save_memory:293 self.model.low_vram_shift(is_diffusing=True)294 295 samples, intermediates = self.ddim_sampler.sample(296 ddim_steps,297 num_samples,298 shape,299 cond,300 verbose=False,301 eta=eta,302 unconditional_guidance_scale=scale,303 unconditional_conditioning=un_cond)304 305 if config.save_memory:306 self.model.low_vram_shift(is_diffusing=False)307 308 x_samples = self.model.decode_first_stage(samples)309 x_samples = (310 einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +311 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)312 313 results = [x_samples[i] for i in range(num_samples)]314 return [detected_map] + results315 316 @torch.inference_mode()317 def process_scribble(self, input_image, prompt, a_prompt, n_prompt,318 num_samples, image_resolution, ddim_steps, scale,319 seed, eta):320 self.load_weight('scribble')321 322 img = resize_image(HWC3(input_image), image_resolution)323 H, W, C = img.shape324 325 detected_map = np.zeros_like(img, dtype=np.uint8)326 detected_map[np.min(img, axis=2) < 127] = 255327 328 control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0329 control = torch.stack([control for _ in range(num_samples)], dim=0)330 control = einops.rearrange(control, 'b h w c -> b c h w').clone()331 332 if seed == -1:333 seed = random.randint(0, 65535)334 seed_everything(seed)335 336 if config.save_memory:337 self.model.low_vram_shift(is_diffusing=False)338 339 cond = {340 'c_concat': [control],341 'c_crossattn': [342 self.model.get_learned_conditioning(343 [prompt + ', ' + a_prompt] * num_samples)344 ]345 }346 un_cond = {347 'c_concat': [control],348 'c_crossattn':349 [self.model.get_learned_conditioning([n_prompt] * num_samples)]350 }351 shape = (4, H // 8, W // 8)352 353 if config.save_memory:354 self.model.low_vram_shift(is_diffusing=True)355 356 samples, intermediates = self.ddim_sampler.sample(357 ddim_steps,358 num_samples,359 shape,360 cond,361 verbose=False,362 eta=eta,363 unconditional_guidance_scale=scale,364 unconditional_conditioning=un_cond)365 366 if config.save_memory:367 self.model.low_vram_shift(is_diffusing=False)368 369 x_samples = self.model.decode_first_stage(samples)370 x_samples = (371 einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +372 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)373 374 results = [x_samples[i] for i in range(num_samples)]375 return [255 - detected_map] + results376 377 @torch.inference_mode()378 def process_scribble_interactive(self, input_image, prompt, a_prompt,379 n_prompt, num_samples, image_resolution,380 ddim_steps, scale, seed, eta):381 self.load_weight('scribble')382 383 img = resize_image(HWC3(input_image['mask'][:, :, 0]),384 image_resolution)385 H, W, C = img.shape386 387 detected_map = np.zeros_like(img, dtype=np.uint8)388 detected_map[np.min(img, axis=2) > 127] = 255389 390 control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0391 control = torch.stack([control for _ in range(num_samples)], dim=0)392 control = einops.rearrange(control, 'b h w c -> b c h w').clone()393 394 if seed == -1:395 seed = random.randint(0, 65535)396 seed_everything(seed)397 398 if config.save_memory:399 self.model.low_vram_shift(is_diffusing=False)400 401 cond = {402 'c_concat': [control],403 'c_crossattn': [404 self.model.get_learned_conditioning(405 [prompt + ', ' + a_prompt] * num_samples)406 ]407 }408 un_cond = {409 'c_concat': [control],410 'c_crossattn':411 [self.model.get_learned_conditioning([n_prompt] * num_samples)]412 }413 shape = (4, H // 8, W // 8)414 415 if config.save_memory:416 self.model.low_vram_shift(is_diffusing=True)417 418 samples, intermediates = self.ddim_sampler.sample(419 ddim_steps,420 num_samples,421 shape,422 cond,423 verbose=False,424 eta=eta,425 unconditional_guidance_scale=scale,426 unconditional_conditioning=un_cond)427 428 if config.save_memory:429 self.model.low_vram_shift(is_diffusing=False)430 431 x_samples = self.model.decode_first_stage(samples)432 x_samples = (433 einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +434 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)435 436 results = [x_samples[i] for i in range(num_samples)]437 return [255 - detected_map] + results438 439 @torch.inference_mode()440 def process_fake_scribble(self, input_image, prompt, a_prompt, n_prompt,441 num_samples, image_resolution, detect_resolution,442 ddim_steps, scale, seed, eta):443 self.load_weight('scribble')444 445 input_image = HWC3(input_image)446 detected_map = apply_hed(resize_image(input_image, detect_resolution))447 detected_map = HWC3(detected_map)448 img = resize_image(input_image, image_resolution)449 H, W, C = img.shape450 451 detected_map = cv2.resize(detected_map, (W, H),452 interpolation=cv2.INTER_LINEAR)453 detected_map = nms(detected_map, 127, 3.0)454 detected_map = cv2.GaussianBlur(detected_map, (0, 0), 3.0)455 detected_map[detected_map > 4] = 255456 detected_map[detected_map < 255] = 0457 458 control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0459 control = torch.stack([control for _ in range(num_samples)], dim=0)460 control = einops.rearrange(control, 'b h w c -> b c h w').clone()461 462 if seed == -1:463 seed = random.randint(0, 65535)464 seed_everything(seed)465 466 if config.save_memory:467 self.model.low_vram_shift(is_diffusing=False)468 469 cond = {470 'c_concat': [control],471 'c_crossattn': [472 self.model.get_learned_conditioning(473 [prompt + ', ' + a_prompt] * num_samples)474 ]475 }476 un_cond = {477 'c_concat': [control],478 'c_crossattn':479 [self.model.get_learned_conditioning([n_prompt] * num_samples)]480 }481 shape = (4, H // 8, W // 8)482 483 if config.save_memory:484 self.model.low_vram_shift(is_diffusing=True)485 486 samples, intermediates = self.ddim_sampler.sample(487 ddim_steps,488 num_samples,489 shape,490 cond,491 verbose=False,492 eta=eta,493 unconditional_guidance_scale=scale,494 unconditional_conditioning=un_cond)495 496 if config.save_memory:497 self.model.low_vram_shift(is_diffusing=False)498 499 x_samples = self.model.decode_first_stage(samples)500 x_samples = (501 einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +502 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)503 504 results = [x_samples[i] for i in range(num_samples)]505 return [255 - detected_map] + results506 507 @torch.inference_mode()508 def process_pose(self, input_image, prompt, a_prompt, n_prompt,509 num_samples, image_resolution, detect_resolution,510 ddim_steps, scale, seed, eta):511 self.load_weight('pose')512 513 input_image = HWC3(input_image)514 detected_map, _ = apply_openpose(515 resize_image(input_image, detect_resolution))516 detected_map = HWC3(detected_map)517 img = resize_image(input_image, image_resolution)518 H, W, C = img.shape519 520 detected_map = cv2.resize(detected_map, (W, H),521 interpolation=cv2.INTER_NEAREST)522 523 control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0524 control = torch.stack([control for _ in range(num_samples)], dim=0)525 control = einops.rearrange(control, 'b h w c -> b c h w').clone()526 527 if seed == -1:528 seed = random.randint(0, 65535)529 seed_everything(seed)530 531 if config.save_memory:532 self.model.low_vram_shift(is_diffusing=False)533 534 cond = {535 'c_concat': [control],536 'c_crossattn': [537 self.model.get_learned_conditioning(538 [prompt + ', ' + a_prompt] * num_samples)539 ]540 }541 un_cond = {542 'c_concat': [control],543 'c_crossattn':544 [self.model.get_learned_conditioning([n_prompt] * num_samples)]545 }546 shape = (4, H // 8, W // 8)547 548 if config.save_memory:549 self.model.low_vram_shift(is_diffusing=True)550 551 samples, intermediates = self.ddim_sampler.sample(552 ddim_steps,553 num_samples,554 shape,555 cond,556 verbose=False,557 eta=eta,558 unconditional_guidance_scale=scale,559 unconditional_conditioning=un_cond)560 561 if config.save_memory:562 self.model.low_vram_shift(is_diffusing=False)563 564 x_samples = self.model.decode_first_stage(samples)565 x_samples = (566 einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +567 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)568 569 results = [x_samples[i] for i in range(num_samples)]570 return [detected_map] + results571 572 @torch.inference_mode()573 def process_seg(self, input_image, prompt, a_prompt, n_prompt, num_samples,574 image_resolution, detect_resolution, ddim_steps, scale,575 seed, eta):576 self.load_weight('seg')577 578 input_image = HWC3(input_image)579 detected_map = apply_uniformer(580 resize_image(input_image, detect_resolution))581 img = resize_image(input_image, image_resolution)582 H, W, C = img.shape583 584 detected_map = cv2.resize(detected_map, (W, H),585 interpolation=cv2.INTER_NEAREST)586 587 control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0588 control = torch.stack([control for _ in range(num_samples)], dim=0)589 control = einops.rearrange(control, 'b h w c -> b c h w').clone()590 591 if seed == -1:592 seed = random.randint(0, 65535)593 seed_everything(seed)594 595 if config.save_memory:596 self.model.low_vram_shift(is_diffusing=False)597 598 cond = {599 'c_concat': [control],600 'c_crossattn': [601 self.model.get_learned_conditioning(602 [prompt + ', ' + a_prompt] * num_samples)603 ]604 }605 un_cond = {606 'c_concat': [control],607 'c_crossattn':608 [self.model.get_learned_conditioning([n_prompt] * num_samples)]609 }610 shape = (4, H // 8, W // 8)611 612 if config.save_memory:613 self.model.low_vram_shift(is_diffusing=True)614 615 samples, intermediates = self.ddim_sampler.sample(616 ddim_steps,617 num_samples,618 shape,619 cond,620 verbose=False,621 eta=eta,622 unconditional_guidance_scale=scale,623 unconditional_conditioning=un_cond)624 625 if config.save_memory:626 self.model.low_vram_shift(is_diffusing=False)627 628 x_samples = self.model.decode_first_stage(samples)629 x_samples = (630 einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +631 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)632 633 results = [x_samples[i] for i in range(num_samples)]634 return [detected_map] + results635 636 @torch.inference_mode()637 def process_depth(self, input_image, prompt, a_prompt, n_prompt,638 num_samples, image_resolution, detect_resolution,639 ddim_steps, scale, seed, eta):640 self.load_weight('depth')641 642 input_image = HWC3(input_image)643 detected_map, _ = apply_midas(644 resize_image(input_image, detect_resolution))645 detected_map = HWC3(detected_map)646 img = resize_image(input_image, image_resolution)647 H, W, C = img.shape648 649 detected_map = cv2.resize(detected_map, (W, H),650 interpolation=cv2.INTER_LINEAR)651 652 control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0653 control = torch.stack([control for _ in range(num_samples)], dim=0)654 control = einops.rearrange(control, 'b h w c -> b c h w').clone()655 656 if seed == -1:657 seed = random.randint(0, 65535)658 seed_everything(seed)659 660 if config.save_memory:661 self.model.low_vram_shift(is_diffusing=False)662 663 cond = {664 'c_concat': [control],665 'c_crossattn': [666 self.model.get_learned_conditioning(667 [prompt + ', ' + a_prompt] * num_samples)668 ]669 }670 un_cond = {671 'c_concat': [control],672 'c_crossattn':673 [self.model.get_learned_conditioning([n_prompt] * num_samples)]674 }675 shape = (4, H // 8, W // 8)676 677 if config.save_memory:678 self.model.low_vram_shift(is_diffusing=True)679 680 samples, intermediates = self.ddim_sampler.sample(681 ddim_steps,682 num_samples,683 shape,684 cond,685 verbose=False,686 eta=eta,687 unconditional_guidance_scale=scale,688 unconditional_conditioning=un_cond)689 690 if config.save_memory:691 self.model.low_vram_shift(is_diffusing=False)692 693 x_samples = self.model.decode_first_stage(samples)694 x_samples = (695 einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +696 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)697 698 results = [x_samples[i] for i in range(num_samples)]699 return [detected_map] + results700 701 @torch.inference_mode()702 def process_normal(self, input_image, prompt, a_prompt, n_prompt,703 num_samples, image_resolution, detect_resolution,704 ddim_steps, scale, seed, eta, bg_threshold):705 self.load_weight('normal')706 707 input_image = HWC3(input_image)708 _, detected_map = apply_midas(resize_image(input_image,709 detect_resolution),710 bg_th=bg_threshold)711 detected_map = HWC3(detected_map)712 img = resize_image(input_image, image_resolution)713 H, W, C = img.shape714 715 detected_map = cv2.resize(detected_map, (W, H),716 interpolation=cv2.INTER_LINEAR)717 718 control = torch.from_numpy(719 detected_map[:, :, ::-1].copy()).float().cuda() / 255.0720 control = torch.stack([control for _ in range(num_samples)], dim=0)721 control = einops.rearrange(control, 'b h w c -> b c h w').clone()722 723 if seed == -1:724 seed = random.randint(0, 65535)725 seed_everything(seed)726 727 if config.save_memory:728 self.model.low_vram_shift(is_diffusing=False)729 730 cond = {731 'c_concat': [control],732 'c_crossattn': [733 self.model.get_learned_conditioning(734 [prompt + ', ' + a_prompt] * num_samples)735 ]736 }737 un_cond = {738 'c_concat': [control],739 'c_crossattn':740 [self.model.get_learned_conditioning([n_prompt] * num_samples)]741 }742 shape = (4, H // 8, W // 8)743 744 if config.save_memory:745 self.model.low_vram_shift(is_diffusing=True)746 747 samples, intermediates = self.ddim_sampler.sample(748 ddim_steps,749 num_samples,750 shape,751 cond,752 verbose=False,753 eta=eta,754 unconditional_guidance_scale=scale,755 unconditional_conditioning=un_cond)756 757 if config.save_memory:758 self.model.low_vram_shift(is_diffusing=False)759 760 x_samples = self.model.decode_first_stage(samples)761 x_samples = (762 einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +763 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)764 765 results = [x_samples[i] for i in range(num_samples)]766 return [detected_map] + results767