quantumcontrol/stable-video-diffusion
1
1import copy2import math3import os4from glob import glob5from typing import Dict, List, Optional, Tuple, Union6 7import cv28import numpy as np9import streamlit as st10import torch11import torch.nn as nn12import torchvision.transforms as TT13from einops import rearrange, repeat14from imwatermark import WatermarkEncoder15from omegaconf import ListConfig, OmegaConf16from PIL import Image17from safetensors.torch import load_file as load_safetensors18from torch import autocast19from torchvision import transforms20from torchvision.utils import make_grid, save_image21 22from scripts.demo.discretization import (Img2ImgDiscretizationWrapper,23 Txt2NoisyDiscretizationWrapper)24from scripts.util.detection.nsfw_and_watermark_dectection import \25 DeepFloydDataFiltering26from sgm.inference.helpers import embed_watermark27from sgm.modules.diffusionmodules.guiders import (LinearPredictionGuider,28 VanillaCFG)29from sgm.modules.diffusionmodules.sampling import (DPMPP2MSampler,30 DPMPP2SAncestralSampler,31 EulerAncestralSampler,32 EulerEDMSampler,33 HeunEDMSampler,34 LinearMultistepSampler)35from sgm.util import append_dims, default, instantiate_from_config36 37 38@st.cache_resource()39def init_st(version_dict, load_ckpt=True, load_filter=True):40 state = dict()41 if not "model" in state:42 config = version_dict["config"]43 ckpt = version_dict["ckpt"]44 45 config = OmegaConf.load(config)46 model, msg = load_model_from_config(config, ckpt if load_ckpt else None)47 48 state["msg"] = msg49 state["model"] = model50 state["ckpt"] = ckpt if load_ckpt else None51 state["config"] = config52 if load_filter:53 state["filter"] = DeepFloydDataFiltering(verbose=False)54 return state55 56 57def load_model(model):58 model.cuda()59 60 61lowvram_mode = False62 63 64def set_lowvram_mode(mode):65 global lowvram_mode66 lowvram_mode = mode67 68 69def initial_model_load(model):70 global lowvram_mode71 if lowvram_mode:72 model.model.half()73 else:74 model.cuda()75 return model76 77 78def unload_model(model):79 global lowvram_mode80 if lowvram_mode:81 model.cpu()82 torch.cuda.empty_cache()83 84 85def load_model_from_config(config, ckpt=None, verbose=True):86 model = instantiate_from_config(config.model)87 88 if ckpt is not None:89 print(f"Loading model from {ckpt}")90 if ckpt.endswith("ckpt"):91 pl_sd = torch.load(ckpt, map_location="cpu")92 if "global_step" in pl_sd:93 global_step = pl_sd["global_step"]94 st.info(f"loaded ckpt from global step {global_step}")95 print(f"Global Step: {pl_sd['global_step']}")96 sd = pl_sd["state_dict"]97 elif ckpt.endswith("safetensors"):98 sd = load_safetensors(ckpt)99 else:100 raise NotImplementedError101 102 msg = None103 104 m, u = model.load_state_dict(sd, strict=False)105 106 if len(m) > 0 and verbose:107 print("missing keys:")108 print(m)109 if len(u) > 0 and verbose:110 print("unexpected keys:")111 print(u)112 else:113 msg = None114 115 model = initial_model_load(model)116 model.eval()117 return model, msg118 119 120def get_unique_embedder_keys_from_conditioner(conditioner):121 return list(set([x.input_key for x in conditioner.embedders]))122 123 124def init_embedder_options(keys, init_dict, prompt=None, negative_prompt=None):125 # Hardcoded demo settings; might undergo some changes in the future126 127 value_dict = {}128 for key in keys:129 if key == "txt":130 if prompt is None:131 prompt = "A professional photograph of an astronaut riding a pig"132 if negative_prompt is None:133 negative_prompt = ""134 135 prompt = st.text_input("Prompt", prompt)136 negative_prompt = st.text_input("Negative prompt", negative_prompt)137 138 value_dict["prompt"] = prompt139 value_dict["negative_prompt"] = negative_prompt140 141 if key == "original_size_as_tuple":142 orig_width = st.number_input(143 "orig_width",144 value=init_dict["orig_width"],145 min_value=16,146 )147 orig_height = st.number_input(148 "orig_height",149 value=init_dict["orig_height"],150 min_value=16,151 )152 153 value_dict["orig_width"] = orig_width154 value_dict["orig_height"] = orig_height155 156 if key == "crop_coords_top_left":157 crop_coord_top = st.number_input("crop_coords_top", value=0, min_value=0)158 crop_coord_left = st.number_input("crop_coords_left", value=0, min_value=0)159 160 value_dict["crop_coords_top"] = crop_coord_top161 value_dict["crop_coords_left"] = crop_coord_left162 163 if key == "aesthetic_score":164 value_dict["aesthetic_score"] = 6.0165 value_dict["negative_aesthetic_score"] = 2.5166 167 if key == "target_size_as_tuple":168 value_dict["target_width"] = init_dict["target_width"]169 value_dict["target_height"] = init_dict["target_height"]170 171 if key in ["fps_id", "fps"]:172 fps = st.number_input("fps", value=6, min_value=1)173 174 value_dict["fps"] = fps175 value_dict["fps_id"] = fps - 1176 177 if key == "motion_bucket_id":178 mb_id = st.number_input("motion bucket id", 0, 511, value=127)179 value_dict["motion_bucket_id"] = mb_id180 181 if key == "pool_image":182 st.text("Image for pool conditioning")183 image = load_img(184 key="pool_image_input",185 size=224,186 center_crop=True,187 )188 if image is None:189 st.info("Need an image here")190 image = torch.zeros(1, 3, 224, 224)191 value_dict["pool_image"] = image192 193 return value_dict194 195 196def perform_save_locally(save_path, samples):197 os.makedirs(os.path.join(save_path), exist_ok=True)198 base_count = len(os.listdir(os.path.join(save_path)))199 samples = embed_watermark(samples)200 for sample in samples:201 sample = 255.0 * rearrange(sample.cpu().numpy(), "c h w -> h w c")202 Image.fromarray(sample.astype(np.uint8)).save(203 os.path.join(save_path, f"{base_count:09}.png")204 )205 base_count += 1206 207 208def init_save_locally(_dir, init_value: bool = False):209 save_locally = st.sidebar.checkbox("Save images locally", value=init_value)210 if save_locally:211 save_path = st.text_input("Save path", value=os.path.join(_dir, "samples"))212 else:213 save_path = None214 215 return save_locally, save_path216 217 218def get_guider(options, key):219 guider = st.sidebar.selectbox(220 f"Discretization #{key}",221 [222 "VanillaCFG",223 "IdentityGuider",224 "LinearPredictionGuider",225 ],226 options.get("guider", 0),227 )228 229 additional_guider_kwargs = options.pop("additional_guider_kwargs", {})230 231 if guider == "IdentityGuider":232 guider_config = {233 "target": "sgm.modules.diffusionmodules.guiders.IdentityGuider"234 }235 elif guider == "VanillaCFG":236 scale_schedule = st.sidebar.selectbox(237 f"Scale schedule #{key}",238 ["Identity", "Oscillating"],239 )240 241 if scale_schedule == "Identity":242 scale = st.number_input(243 f"cfg-scale #{key}",244 value=options.get("cfg", 5.0),245 min_value=0.0,246 )247 248 scale_schedule_config = {249 "target": "sgm.modules.diffusionmodules.guiders.IdentitySchedule",250 "params": {"scale": scale},251 }252 253 elif scale_schedule == "Oscillating":254 small_scale = st.number_input(255 f"small cfg-scale #{key}",256 value=4.0,257 min_value=0.0,258 )259 260 large_scale = st.number_input(261 f"large cfg-scale #{key}",262 value=16.0,263 min_value=0.0,264 )265 266 sigma_cutoff = st.number_input(267 f"sigma cutoff #{key}",268 value=1.0,269 min_value=0.0,270 )271 272 scale_schedule_config = {273 "target": "sgm.modules.diffusionmodules.guiders.OscillatingSchedule",274 "params": {275 "small_scale": small_scale,276 "large_scale": large_scale,277 "sigma_cutoff": sigma_cutoff,278 },279 }280 else:281 raise NotImplementedError282 283 guider_config = {284 "target": "sgm.modules.diffusionmodules.guiders.VanillaCFG",285 "params": {286 "scale_schedule_config": scale_schedule_config,287 **additional_guider_kwargs,288 },289 }290 elif guider == "LinearPredictionGuider":291 max_scale = st.number_input(292 f"max-cfg-scale #{key}",293 value=options.get("cfg", 1.5),294 min_value=1.0,295 )296 min_scale = st.number_input(297 f"min guidance scale",298 value=options.get("min_cfg", 1.0),299 min_value=1.0,300 max_value=10.0,301 )302 303 guider_config = {304 "target": "sgm.modules.diffusionmodules.guiders.LinearPredictionGuider",305 "params": {306 "max_scale": max_scale,307 "min_scale": min_scale,308 "num_frames": options["num_frames"],309 **additional_guider_kwargs,310 },311 }312 else:313 raise NotImplementedError314 return guider_config315 316 317def init_sampling(318 key=1,319 img2img_strength: Optional[float] = None,320 specify_num_samples: bool = True,321 stage2strength: Optional[float] = None,322 options: Optional[Dict[str, int]] = None,323):324 options = {} if options is None else options325 326 num_rows, num_cols = 1, 1327 if specify_num_samples:328 num_cols = st.number_input(329 f"num cols #{key}", value=num_cols, min_value=1, max_value=10330 )331 332 steps = st.sidebar.number_input(333 f"steps #{key}", value=options.get("num_steps", 40), min_value=1, max_value=1000334 )335 sampler = st.sidebar.selectbox(336 f"Sampler #{key}",337 [338 "EulerEDMSampler",339 "HeunEDMSampler",340 "EulerAncestralSampler",341 "DPMPP2SAncestralSampler",342 "DPMPP2MSampler",343 "LinearMultistepSampler",344 ],345 options.get("sampler", 0),346 )347 discretization = st.sidebar.selectbox(348 f"Discretization #{key}",349 [350 "LegacyDDPMDiscretization",351 "EDMDiscretization",352 ],353 options.get("discretization", 0),354 )355 356 discretization_config = get_discretization(discretization, options=options, key=key)357 358 guider_config = get_guider(options=options, key=key)359 360 sampler = get_sampler(sampler, steps, discretization_config, guider_config, key=key)361 if img2img_strength is not None:362 st.warning(363 f"Wrapping {sampler.__class__.__name__} with Img2ImgDiscretizationWrapper"364 )365 sampler.discretization = Img2ImgDiscretizationWrapper(366 sampler.discretization, strength=img2img_strength367 )368 if stage2strength is not None:369 sampler.discretization = Txt2NoisyDiscretizationWrapper(370 sampler.discretization, strength=stage2strength, original_steps=steps371 )372 return sampler, num_rows, num_cols373 374 375def get_discretization(discretization, options, key=1):376 if discretization == "LegacyDDPMDiscretization":377 discretization_config = {378 "target": "sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization",379 }380 elif discretization == "EDMDiscretization":381 sigma_min = st.number_input(382 f"sigma_min #{key}", value=options.get("sigma_min", 0.03)383 ) # 0.0292384 sigma_max = st.number_input(385 f"sigma_max #{key}", value=options.get("sigma_max", 14.61)386 ) # 14.6146387 rho = st.number_input(f"rho #{key}", value=options.get("rho", 3.0))388 discretization_config = {389 "target": "sgm.modules.diffusionmodules.discretizer.EDMDiscretization",390 "params": {391 "sigma_min": sigma_min,392 "sigma_max": sigma_max,393 "rho": rho,394 },395 }396 397 return discretization_config398 399 400def get_sampler(sampler_name, steps, discretization_config, guider_config, key=1):401 if sampler_name == "EulerEDMSampler" or sampler_name == "HeunEDMSampler":402 s_churn = st.sidebar.number_input(f"s_churn #{key}", value=0.0, min_value=0.0)403 s_tmin = st.sidebar.number_input(f"s_tmin #{key}", value=0.0, min_value=0.0)404 s_tmax = st.sidebar.number_input(f"s_tmax #{key}", value=999.0, min_value=0.0)405 s_noise = st.sidebar.number_input(f"s_noise #{key}", value=1.0, min_value=0.0)406 407 if sampler_name == "EulerEDMSampler":408 sampler = EulerEDMSampler(409 num_steps=steps,410 discretization_config=discretization_config,411 guider_config=guider_config,412 s_churn=s_churn,413 s_tmin=s_tmin,414 s_tmax=s_tmax,415 s_noise=s_noise,416 verbose=True,417 )418 elif sampler_name == "HeunEDMSampler":419 sampler = HeunEDMSampler(420 num_steps=steps,421 discretization_config=discretization_config,422 guider_config=guider_config,423 s_churn=s_churn,424 s_tmin=s_tmin,425 s_tmax=s_tmax,426 s_noise=s_noise,427 verbose=True,428 )429 elif (430 sampler_name == "EulerAncestralSampler"431 or sampler_name == "DPMPP2SAncestralSampler"432 ):433 s_noise = st.sidebar.number_input("s_noise", value=1.0, min_value=0.0)434 eta = st.sidebar.number_input("eta", value=1.0, min_value=0.0)435 436 if sampler_name == "EulerAncestralSampler":437 sampler = EulerAncestralSampler(438 num_steps=steps,439 discretization_config=discretization_config,440 guider_config=guider_config,441 eta=eta,442 s_noise=s_noise,443 verbose=True,444 )445 elif sampler_name == "DPMPP2SAncestralSampler":446 sampler = DPMPP2SAncestralSampler(447 num_steps=steps,448 discretization_config=discretization_config,449 guider_config=guider_config,450 eta=eta,451 s_noise=s_noise,452 verbose=True,453 )454 elif sampler_name == "DPMPP2MSampler":455 sampler = DPMPP2MSampler(456 num_steps=steps,457 discretization_config=discretization_config,458 guider_config=guider_config,459 verbose=True,460 )461 elif sampler_name == "LinearMultistepSampler":462 order = st.sidebar.number_input("order", value=4, min_value=1)463 sampler = LinearMultistepSampler(464 num_steps=steps,465 discretization_config=discretization_config,466 guider_config=guider_config,467 order=order,468 verbose=True,469 )470 else:471 raise ValueError(f"unknown sampler {sampler_name}!")472 473 return sampler474 475 476def get_interactive_image() -> Image.Image:477 image = st.file_uploader("Input", type=["jpg", "JPEG", "png"])478 if image is not None:479 image = Image.open(image)480 if not image.mode == "RGB":481 image = image.convert("RGB")482 return image483 484 485def load_img(486 display: bool = True,487 size: Union[None, int, Tuple[int, int]] = None,488 center_crop: bool = False,489):490 image = get_interactive_image()491 if image is None:492 return None493 if display:494 st.image(image)495 w, h = image.size496 print(f"loaded input image of size ({w}, {h})")497 498 transform = []499 if size is not None:500 transform.append(transforms.Resize(size))501 if center_crop:502 transform.append(transforms.CenterCrop(size))503 transform.append(transforms.ToTensor())504 transform.append(transforms.Lambda(lambda x: 2.0 * x - 1.0))505 506 transform = transforms.Compose(transform)507 img = transform(image)[None, ...]508 st.text(f"input min/max/mean: {img.min():.3f}/{img.max():.3f}/{img.mean():.3f}")509 return img510 511 512def get_init_img(batch_size=1, key=None):513 init_image = load_img(key=key).cuda()514 init_image = repeat(init_image, "1 ... -> b ...", b=batch_size)515 return init_image516 517 518def do_sample(519 model,520 sampler,521 value_dict,522 num_samples,523 H,524 W,525 C,526 F,527 force_uc_zero_embeddings: Optional[List] = None,528 force_cond_zero_embeddings: Optional[List] = None,529 batch2model_input: List = None,530 return_latents=False,531 filter=None,532 T=None,533 additional_batch_uc_fields=None,534 decoding_t=None,535):536 force_uc_zero_embeddings = default(force_uc_zero_embeddings, [])537 batch2model_input = default(batch2model_input, [])538 additional_batch_uc_fields = default(additional_batch_uc_fields, [])539 540 st.text("Sampling")541 542 outputs = st.empty()543 precision_scope = autocast544 with torch.no_grad():545 with precision_scope("cuda"):546 with model.ema_scope():547 if T is not None:548 num_samples = [num_samples, T]549 else:550 num_samples = [num_samples]551 552 load_model(model.conditioner)553 batch, batch_uc = get_batch(554 get_unique_embedder_keys_from_conditioner(model.conditioner),555 value_dict,556 num_samples,557 T=T,558 additional_batch_uc_fields=additional_batch_uc_fields,559 )560 561 c, uc = model.conditioner.get_unconditional_conditioning(562 batch,563 batch_uc=batch_uc,564 force_uc_zero_embeddings=force_uc_zero_embeddings,565 force_cond_zero_embeddings=force_cond_zero_embeddings,566 )567 unload_model(model.conditioner)568 569 for k in c:570 if not k == "crossattn":571 c[k], uc[k] = map(572 lambda y: y[k][: math.prod(num_samples)].to("cuda"), (c, uc)573 )574 if k in ["crossattn", "concat"] and T is not None:575 uc[k] = repeat(uc[k], "b ... -> b t ...", t=T)576 uc[k] = rearrange(uc[k], "b t ... -> (b t) ...", t=T)577 c[k] = repeat(c[k], "b ... -> b t ...", t=T)578 c[k] = rearrange(c[k], "b t ... -> (b t) ...", t=T)579 580 additional_model_inputs = {}581 for k in batch2model_input:582 if k == "image_only_indicator":583 assert T is not None584 585 if isinstance(586 sampler.guider, (VanillaCFG, LinearPredictionGuider)587 ):588 additional_model_inputs[k] = torch.zeros(589 num_samples[0] * 2, num_samples[1]590 ).to("cuda")591 else:592 additional_model_inputs[k] = torch.zeros(num_samples).to(593 "cuda"594 )595 else:596 additional_model_inputs[k] = batch[k]597 598 shape = (math.prod(num_samples), C, H // F, W // F)599 randn = torch.randn(shape).to("cuda")600 601 def denoiser(input, sigma, c):602 return model.denoiser(603 model.model, input, sigma, c, **additional_model_inputs604 )605 606 load_model(model.denoiser)607 load_model(model.model)608 samples_z = sampler(denoiser, randn, cond=c, uc=uc)609 unload_model(model.model)610 unload_model(model.denoiser)611 612 load_model(model.first_stage_model)613 model.en_and_decode_n_samples_a_time = (614 decoding_t # Decode n frames at a time615 )616 samples_x = model.decode_first_stage(samples_z)617 samples = torch.clamp((samples_x + 1.0) / 2.0, min=0.0, max=1.0)618 unload_model(model.first_stage_model)619 620 if filter is not None:621 samples = filter(samples)622 623 if T is None:624 grid = torch.stack([samples])625 grid = rearrange(grid, "n b c h w -> (n h) (b w) c")626 outputs.image(grid.cpu().numpy())627 else:628 as_vids = rearrange(samples, "(b t) c h w -> b t c h w", t=T)629 for i, vid in enumerate(as_vids):630 grid = rearrange(make_grid(vid, nrow=4), "c h w -> h w c")631 st.image(632 grid.cpu().numpy(),633 f"Sample #{i} as image",634 )635 636 if return_latents:637 return samples, samples_z638 return samples639 640 641def get_batch(642 keys,643 value_dict: dict,644 N: Union[List, ListConfig],645 device: str = "cuda",646 T: int = None,647 additional_batch_uc_fields: List[str] = [],648):649 # Hardcoded demo setups; might undergo some changes in the future650 651 batch = {}652 batch_uc = {}653 654 for key in keys:655 if key == "txt":656 batch["txt"] = [value_dict["prompt"]] * math.prod(N)657 658 batch_uc["txt"] = [value_dict["negative_prompt"]] * math.prod(N)659 660 elif key == "original_size_as_tuple":661 batch["original_size_as_tuple"] = (662 torch.tensor([value_dict["orig_height"], value_dict["orig_width"]])663 .to(device)664 .repeat(math.prod(N), 1)665 )666 elif key == "crop_coords_top_left":667 batch["crop_coords_top_left"] = (668 torch.tensor(669 [value_dict["crop_coords_top"], value_dict["crop_coords_left"]]670 )671 .to(device)672 .repeat(math.prod(N), 1)673 )674 elif key == "aesthetic_score":675 batch["aesthetic_score"] = (676 torch.tensor([value_dict["aesthetic_score"]])677 .to(device)678 .repeat(math.prod(N), 1)679 )680 batch_uc["aesthetic_score"] = (681 torch.tensor([value_dict["negative_aesthetic_score"]])682 .to(device)683 .repeat(math.prod(N), 1)684 )685 686 elif key == "target_size_as_tuple":687 batch["target_size_as_tuple"] = (688 torch.tensor([value_dict["target_height"], value_dict["target_width"]])689 .to(device)690 .repeat(math.prod(N), 1)691 )692 elif key == "fps":693 batch[key] = (694 torch.tensor([value_dict["fps"]]).to(device).repeat(math.prod(N))695 )696 elif key == "fps_id":697 batch[key] = (698 torch.tensor([value_dict["fps_id"]]).to(device).repeat(math.prod(N))699 )700 elif key == "motion_bucket_id":701 batch[key] = (702 torch.tensor([value_dict["motion_bucket_id"]])703 .to(device)704 .repeat(math.prod(N))705 )706 elif key == "pool_image":707 batch[key] = repeat(value_dict[key], "1 ... -> b ...", b=math.prod(N)).to(708 device, dtype=torch.half709 )710 elif key == "cond_aug":711 batch[key] = repeat(712 torch.tensor([value_dict["cond_aug"]]).to("cuda"),713 "1 -> b",714 b=math.prod(N),715 )716 elif key == "cond_frames":717 batch[key] = repeat(value_dict["cond_frames"], "1 ... -> b ...", b=N[0])718 elif key == "cond_frames_without_noise":719 batch[key] = repeat(720 value_dict["cond_frames_without_noise"], "1 ... -> b ...", b=N[0]721 )722 else:723 batch[key] = value_dict[key]724 725 if T is not None:726 batch["num_video_frames"] = T727 728 for key in batch.keys():729 if key not in batch_uc and isinstance(batch[key], torch.Tensor):730 batch_uc[key] = torch.clone(batch[key])731 elif key in additional_batch_uc_fields and key not in batch_uc:732 batch_uc[key] = copy.copy(batch[key])733 return batch, batch_uc734 735 736@torch.no_grad()737def do_img2img(738 img,739 model,740 sampler,741 value_dict,742 num_samples,743 force_uc_zero_embeddings: Optional[List] = None,744 force_cond_zero_embeddings: Optional[List] = None,745 additional_kwargs={},746 offset_noise_level: int = 0.0,747 return_latents=False,748 skip_encode=False,749 filter=None,750 add_noise=True,751):752 st.text("Sampling")753 754 outputs = st.empty()755 precision_scope = autocast756 with torch.no_grad():757 with precision_scope("cuda"):758 with model.ema_scope():759 load_model(model.conditioner)760 batch, batch_uc = get_batch(761 get_unique_embedder_keys_from_conditioner(model.conditioner),762 value_dict,763 [num_samples],764 )765 c, uc = model.conditioner.get_unconditional_conditioning(766 batch,767 batch_uc=batch_uc,768 force_uc_zero_embeddings=force_uc_zero_embeddings,769 force_cond_zero_embeddings=force_cond_zero_embeddings,770 )771 unload_model(model.conditioner)772 for k in c:773 c[k], uc[k] = map(lambda y: y[k][:num_samples].to("cuda"), (c, uc))774 775 for k in additional_kwargs:776 c[k] = uc[k] = additional_kwargs[k]777 if skip_encode:778 z = img779 else:780 load_model(model.first_stage_model)781 z = model.encode_first_stage(img)782 unload_model(model.first_stage_model)783 784 noise = torch.randn_like(z)785 786 sigmas = sampler.discretization(sampler.num_steps).cuda()787 sigma = sigmas[0]788 789 st.info(f"all sigmas: {sigmas}")790 st.info(f"noising sigma: {sigma}")791 if offset_noise_level > 0.0:792 noise = noise + offset_noise_level * append_dims(793 torch.randn(z.shape[0], device=z.device), z.ndim794 )795 if add_noise:796 noised_z = z + noise * append_dims(sigma, z.ndim).cuda()797 noised_z = noised_z / torch.sqrt(798 1.0 + sigmas[0] ** 2.0799 ) # Note: hardcoded to DDPM-like scaling. need to generalize later.800 else:801 noised_z = z / torch.sqrt(1.0 + sigmas[0] ** 2.0)802 803 def denoiser(x, sigma, c):804 return model.denoiser(model.model, x, sigma, c)805 806 load_model(model.denoiser)807 load_model(model.model)808 samples_z = sampler(denoiser, noised_z, cond=c, uc=uc)809 unload_model(model.model)810 unload_model(model.denoiser)811 812 load_model(model.first_stage_model)813 samples_x = model.decode_first_stage(samples_z)814 unload_model(model.first_stage_model)815 samples = torch.clamp((samples_x + 1.0) / 2.0, min=0.0, max=1.0)816 817 if filter is not None:818 samples = filter(samples)819 820 grid = rearrange(grid, "n b c h w -> (n h) (b w) c")821 outputs.image(grid.cpu().numpy())822 if return_latents:823 return samples, samples_z824 return samples825 826 827def get_resizing_factor(828 desired_shape: Tuple[int, int], current_shape: Tuple[int, int]829) -> float:830 r_bound = desired_shape[1] / desired_shape[0]831 aspect_r = current_shape[1] / current_shape[0]832 if r_bound >= 1.0:833 if aspect_r >= r_bound:834 factor = min(desired_shape) / min(current_shape)835 else:836 if aspect_r < 1.0:837 factor = max(desired_shape) / min(current_shape)838 else:839 factor = max(desired_shape) / max(current_shape)840 else:841 if aspect_r <= r_bound:842 factor = min(desired_shape) / min(current_shape)843 else:844 if aspect_r > 1:845 factor = max(desired_shape) / min(current_shape)846 else:847 factor = max(desired_shape) / max(current_shape)848 849 return factor850 851 852def get_interactive_image(key=None) -> Image.Image:853 image = st.file_uploader("Input", type=["jpg", "JPEG", "png"], key=key)854 if image is not None:855 image = Image.open(image)856 if not image.mode == "RGB":857 image = image.convert("RGB")858 return image859 860 861def load_img_for_prediction(862 W: int, H: int, display=True, key=None, device="cuda"863) -> torch.Tensor:864 image = get_interactive_image(key=key)865 if image is None:866 return None867 if display:868 st.image(image)869 w, h = image.size870 871 image = np.array(image).transpose(2, 0, 1)872 image = torch.from_numpy(image).to(dtype=torch.float32) / 255.0873 image = image.unsqueeze(0)874 875 rfs = get_resizing_factor((H, W), (h, w))876 resize_size = [int(np.ceil(rfs * s)) for s in (h, w)]877 top = (resize_size[0] - H) // 2878 left = (resize_size[1] - W) // 2879 880 image = torch.nn.functional.interpolate(881 image, resize_size, mode="area", antialias=False882 )883 image = TT.functional.crop(image, top=top, left=left, height=H, width=W)884 885 if display:886 numpy_img = np.transpose(image[0].numpy(), (1, 2, 0))887 pil_image = Image.fromarray((numpy_img * 255).astype(np.uint8))888 st.image(pil_image)889 return image.to(device) * 2.0 - 1.0890 891 892def save_video_as_grid_and_mp4(893 video_batch: torch.Tensor, save_path: str, T: int, fps: int = 5894):895 os.makedirs(save_path, exist_ok=True)896 base_count = len(glob(os.path.join(save_path, "*.mp4")))897 898 video_batch = rearrange(video_batch, "(b t) c h w -> b t c h w", t=T)899 video_batch = embed_watermark(video_batch)900 for vid in video_batch:901 save_image(vid, fp=os.path.join(save_path, f"{base_count:06d}.png"), nrow=4)902 903 video_path = os.path.join(save_path, f"{base_count:06d}.mp4")904 905 writer = cv2.VideoWriter(906 video_path,907 cv2.VideoWriter_fourcc(*"MP4V"),908 fps,909 (vid.shape[-1], vid.shape[-2]),910 )911 912 vid = (913 (rearrange(vid, "t c h w -> t h w c") * 255).cpu().numpy().astype(np.uint8)914 )915 for frame in vid:916 frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)917 writer.write(frame)918 919 writer.release()920 921 video_path_h264 = video_path[:-4] + "_h264.mp4"922 os.system(f"ffmpeg -i {video_path} -c:v libx264 {video_path_h264}")923 924 with open(video_path_h264, "rb") as f:925 video_bytes = f.read()926 st.video(video_bytes)927 928 base_count += 1929