from transformers import CLIPTextModel, CLIPTokenizer, logging from diffusers import AutoencoderKL, UNet2DConditionModel, DDIMScheduler from diffusers.utils.torch_utils import randn_tensor # suppress partial model loading warning logging.set_verbosity_error() import os from tqdm import tqdm, trange import torch import torch.nn as nn import argparse from torchvision.io import write_video from pathlib import Path from utils import * import torchvision.transforms as T import cv2 import numpy as np def get_timesteps(scheduler, num_inference_steps, strength, device): # get the original timestep using init_timestep init_timestep = min(int(num_inference_steps * strength), num_inference_steps) t_start = max(num_inference_steps - init_timestep, 0) timesteps = scheduler.timesteps[t_start:] return timesteps, num_inference_steps - t_start class Preprocess(nn.Module): def __init__(self, device, opt, vae, tokenizer, text_encoder, unet,scheduler, hf_key=None): super().__init__() self.device = device self.sd_version = opt["sd_version"] self.use_depth = False self.config = opt print(f'[INFO] loading stable diffusion...') if hf_key is not None: print(f'[INFO] using hugging face custom model key: {hf_key}') model_key = hf_key elif self.sd_version == '2.1': model_key = "stabilityai/stable-diffusion-2-1-base" elif self.sd_version == '2.0': model_key = "stabilityai/stable-diffusion-2-base" elif self.sd_version == '1.5' or self.sd_version == 'ControlNet': model_key = "runwayml/stable-diffusion-v1-5" elif self.sd_version == 'depth': model_key = "stabilityai/stable-diffusion-2-depth" else: raise ValueError(f'Stable-diffusion version {self.sd_version} not supported.') self.model_key = model_key # Create model # self.vae = AutoencoderKL.from_pretrained(model_key, subfolder="vae", revision="fp16", # torch_dtype=torch.float16).to(self.device) # self.tokenizer = CLIPTokenizer.from_pretrained(model_key, subfolder="tokenizer") # self.text_encoder = CLIPTextModel.from_pretrained(model_key, subfolder="text_encoder", revision="fp16", # torch_dtype=torch.float16).to(self.device) # self.unet = UNet2DConditionModel.from_pretrained(model_key, subfolder="unet", revision="fp16", # torch_dtype=torch.float16).to(self.device) self.vae = vae self.tokenizer = tokenizer self.text_encoder = text_encoder self.unet = unet self.scheduler=scheduler self.total_inverted_latents = {} self.noise_total = None # will contain all zs if inversion == 'ddpm', var name chosen to match the save path of zs used in pr https://github.com/omerbt/TokenFlow/pull/24/files# self.paths, self.frames, self.latents = self.get_data(self.config["data_path"], self.config["n_frames"]) print("self.frames", self.frames.shape) print("self.latents", self.latents.shape) if self.sd_version == 'ControlNet': from diffusers import ControlNetModel, StableDiffusionControlNetPipeline controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16).to(self.device) control_pipe = StableDiffusionControlNetPipeline.from_pretrained( "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16 ).to(self.device) self.unet = control_pipe.unet self.controlnet = control_pipe.controlnet self.canny_cond = self.get_canny_cond() elif self.sd_version == 'depth': self.depth_maps = self.prepare_depth_maps() self.scheduler = scheduler self.unet.enable_xformers_memory_efficient_attention() print(f'[INFO] loaded stable diffusion!') @torch.no_grad() def prepare_depth_maps(self, model_type='DPT_Large', device='cuda'): depth_maps = [] midas = torch.hub.load("intel-isl/MiDaS", model_type) midas.to(device) midas.eval() midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms") if model_type == "DPT_Large" or model_type == "DPT_Hybrid": transform = midas_transforms.dpt_transform else: transform = midas_transforms.small_transform for i in range(len(self.paths)): img = cv2.imread(self.paths[i]) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) latent_h = img.shape[0] // 8 latent_w = img.shape[1] // 8 input_batch = transform(img).to(device) prediction = midas(input_batch) depth_map = torch.nn.functional.interpolate( prediction.unsqueeze(1), size=(latent_h, latent_w), mode="bicubic", align_corners=False, ) depth_min = torch.amin(depth_map, dim=[1, 2, 3], keepdim=True) depth_max = torch.amax(depth_map, dim=[1, 2, 3], keepdim=True) depth_map = 2.0 * (depth_map - depth_min) / (depth_max - depth_min) - 1.0 depth_maps.append(depth_map) return torch.cat(depth_maps).to(self.device).to(torch.float16) @torch.no_grad() def get_canny_cond(self): canny_cond = [] for image in self.frames.cpu().permute(0, 2, 3, 1): image = np.uint8(np.array(255 * image)) low_threshold = 100 high_threshold = 200 image = cv2.Canny(image, low_threshold, high_threshold) image = image[:, :, None] image = np.concatenate([image, image, image], axis=2) image = torch.from_numpy((image.astype(np.float32) / 255.0)) canny_cond.append(image) canny_cond = torch.stack(canny_cond).permute(0, 3, 1, 2).to(self.device).to(torch.float16) return canny_cond def controlnet_pred(self, latent_model_input, t, text_embed_input, controlnet_cond): down_block_res_samples, mid_block_res_sample = self.controlnet( latent_model_input, t, encoder_hidden_states=text_embed_input, controlnet_cond=controlnet_cond, conditioning_scale=1, return_dict=False, ) # apply the denoising network noise_pred = self.unet( latent_model_input, t, encoder_hidden_states=text_embed_input, cross_attention_kwargs={}, down_block_additional_residuals=down_block_res_samples, mid_block_additional_residual=mid_block_res_sample, return_dict=False, )[0] return noise_pred @torch.no_grad() def encode_text(self, prompts, device=None): if device is None: device = self.device text_inputs = self.tokenizer( prompts, padding="max_length", max_length=self.tokenizer.model_max_length, return_tensors="pt", ) text_input_ids = text_inputs.input_ids if text_input_ids.shape[-1] > self.tokenizer.model_max_length: removed_text = self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length:]) print( "The following part of your input was truncated because CLIP can only handle sequences up to" f" {self.tokenizer.model_max_length} tokens: {removed_text}" ) text_input_ids = text_input_ids[:, : self.tokenizer.model_max_length] text_embeddings = self.text_encoder(text_input_ids.to(device))[0] return text_embeddings @torch.no_grad() def get_text_embeds(self, prompt, negative_prompt, device="cuda"): text_embeddings = self.encode_text(prompt, device=device) uncond_embeddings = self.encode_text(negative_prompt, device=device) text_embeddings = torch.cat([uncond_embeddings, text_embeddings]) return text_embeddings @torch.no_grad() def decode_latents(self, latents): decoded = [] batch_size = 8 for b in range(0, latents.shape[0], batch_size): latents_batch = 1 / 0.18215 * latents[b:b + batch_size] imgs = self.vae.decode(latents_batch).sample imgs = (imgs / 2 + 0.5).clamp(0, 1) decoded.append(imgs) return torch.cat(decoded) @torch.no_grad() def encode_imgs(self, imgs, batch_size=10, deterministic=True): imgs = 2 * imgs - 1 latents = [] for i in range(0, len(imgs), batch_size): posterior = self.vae.encode(imgs[i:i + batch_size]).latent_dist latent = posterior.mean if deterministic else posterior.sample() latents.append(latent * self.vae.config.scaling_factor) latents = torch.cat(latents) return latents def get_data(self, frames_path, n_frames): # load frames if not self.config["frames"]: paths = [f"{frames_path}/%05d.png" % i for i in range(n_frames)] print(paths) if not os.path.exists(paths[0]): paths = [f"{frames_path}/%05d.jpg" % i for i in range(n_frames)] self.paths = paths frames = [Image.open(path).convert('RGB') for path in paths] if frames[0].size[0] == frames[0].size[1]: frames = [frame.resize((512, 512), resample=Image.Resampling.LANCZOS) for frame in frames] else: frames = self.config["frames"][:n_frames] frames = torch.stack([T.ToTensor()(frame) for frame in frames]).to(torch.float16).to(self.device) # encode to latents latents = self.encode_imgs(frames, deterministic=True).to(torch.float16).to(self.device) print("frames", frames.shape) print("latents", latents.shape) if not self.config["frames"]: return paths, frames, latents else: return None, frames, latents @torch.no_grad() def ddim_inversion(self, cond, latent_frames, save_path, batch_size, save_latents=True, timesteps_to_save=None): timesteps = reversed(self.scheduler.timesteps) timesteps_to_save = timesteps_to_save if timesteps_to_save is not None else timesteps return_inverted_latents = self.config["frames"] is not None for i, t in enumerate(tqdm(timesteps)): for b in range(0, latent_frames.shape[0], int(batch_size)): x_batch = latent_frames[b:b + batch_size] model_input = x_batch cond_batch = cond.repeat(x_batch.shape[0], 1, 1) if self.sd_version == 'depth': depth_maps = torch.cat([self.depth_maps[b: b + batch_size]]) model_input = torch.cat([x_batch, depth_maps],dim=1) alpha_prod_t = self.scheduler.alphas_cumprod[t] alpha_prod_t_prev = ( self.scheduler.alphas_cumprod[timesteps[i - 1]] if i > 0 else self.scheduler.final_alpha_cumprod ) mu = alpha_prod_t ** 0.5 mu_prev = alpha_prod_t_prev ** 0.5 sigma = (1 - alpha_prod_t) ** 0.5 sigma_prev = (1 - alpha_prod_t_prev) ** 0.5 eps = self.unet(model_input, t, encoder_hidden_states=cond_batch).sample if self.sd_version != 'ControlNet' \ else self.controlnet_pred(x_batch, t, cond_batch, torch.cat([self.canny_cond[b: b + batch_size]])) pred_x0 = (x_batch - sigma_prev * eps) / mu_prev latent_frames[b:b + batch_size] = mu * pred_x0 + sigma * eps if return_inverted_latents and t in timesteps_to_save: self.total_inverted_latents[f'noisy_latents_{t}'] = latent_frames.clone() if save_latents and t in timesteps_to_save: torch.save(latent_frames, os.path.join(save_path, 'latents', f'noisy_latents_{t}.pt')) if save_latents: torch.save(latent_frames, os.path.join(save_path, 'latents', f'noisy_latents_{t}.pt')) if return_inverted_latents: self.total_inverted_latents[f'noisy_latents_{t}'] = latent_frames.clone() return latent_frames @torch.no_grad() def ddpm_inversion(self, cond, latent_frames, batch_size, num_inversion_steps, save_path=None, save_latents=True, eta: float = 1.0, skip_steps=20): timesteps = self.scheduler.timesteps return_inverted_latents = self.config["frames"] is not None variance_noise_shape = ( num_inversion_steps, *latent_frames.shape) x0 = latent_frames t_to_idx = {int(v): k for k, v in enumerate(timesteps)} xts = torch.zeros(size=variance_noise_shape, device=self.device, dtype=cond.dtype) for t in reversed(timesteps): idx = t_to_idx[int(t)] for b in range(0, x0.shape[0], batch_size): x_batch = x0[b:b + batch_size] noise = randn_tensor(shape=x_batch.shape, device=self.device, dtype=x0.dtype) xts[idx, b:b + batch_size] = self.scheduler.add_noise(x_batch, noise, t) xts = torch.cat([xts, x0.unsqueeze(0)], dim=0) zs = torch.zeros(size=variance_noise_shape, device=self.device, dtype=cond.dtype) for t in tqdm(timesteps): idx = t_to_idx[int(t)] # 1. predict noise residual for b in range(0, x0.shape[0], batch_size): xt = xts[idx, b:b + batch_size] cond_batch = cond.repeat(xt.shape[0], 1, 1) noise_pred = self.unet(xt, timestep=t, encoder_hidden_states=cond_batch).sample xtm1 = xts[idx + 1, b:b + batch_size] z, xtm1_corrected = compute_noise(self.scheduler, xtm1, xt, t, noise_pred, eta) zs[idx, b:b + batch_size] = z # correction to avoid error accumulation xts[idx + 1, b:b + batch_size] = xtm1_corrected if save_latents: torch.save(xts[idx], os.path.join(save_path, 'latents', f'noisy_latents_{t}.pt')) if return_inverted_latents: self.total_inverted_latents[f'noisy_latents_{t}'] = xts[idx].clone() if save_path: torch.save(xts[idx], os.path.join(save_path, 'latents', f'noisy_latents_{t}.pt')) torch.save(zs, os.path.join(save_path, 'latents', f'noise_total.pt')) if return_inverted_latents: self.total_inverted_latents[f'noisy_latents_{t}'] = xts[idx].clone() self.noise_total = zs.clone() return xts[skip_steps].expand(latent_frames.shape[0], -1, -1, -1), zs def prepare_extra_step_kwargs(self, eta): # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 # and should be between [0, 1] accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) extra_step_kwargs = {} if accepts_eta: extra_step_kwargs["eta"] = eta # check if the scheduler accepts generator accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) return extra_step_kwargs @torch.no_grad() def ddpm_sample(self, init_latents, cond, batch_size, num_inversion_steps, skip_steps, eta, zs_all, guidance_scale=0): use_ddpm = True do_classifier_free_guidance = guidance_scale > 1.0 total_latents = init_latents self.scheduler.set_timesteps(num_inversion_steps, device=device) timesteps = self.scheduler.timesteps zs_total = zs_all[skip_steps:] if use_ddpm: t_to_idx = {int(v): k for k, v in enumerate(timesteps[-zs_total.shape[0]:])} timesteps = timesteps[-zs_total.shape[0]:] num_warmup_steps = len(timesteps) - num_inversion_steps * self.scheduler.order extra_step_kwargs = self.prepare_extra_step_kwargs(eta) for i, t in enumerate(tqdm(timesteps)): for b in range(0, total_latents.shape[0], batch_size): latents = total_latents[b:b + batch_size] if do_classifier_free_guidance: latent_model_input = torch.cat([latents] * 2) else: latent_model_input = latents cond_batch = cond.repeat(latents.shape[0], 1, 1) latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) noise_pred = self.unet( latent_model_input, t, encoder_hidden_states=cond_batch, return_dict=False, )[0] if do_classifier_free_guidance: noise_pred_out = noise_pred.chunk(2) # [b,4, 64, 64] noise_pred_uncond, noise_pred_text = noise_pred_out[0], noise_pred_out[1] # default text guidance noise_guidance = guidance_scale * (noise_pred_text - noise_pred_uncond) noise_pred = noise_pred_uncond + noise_guidance idx = t_to_idx[int(t)] zs = zs_total[idx, b:b + batch_size] latents = self.scheduler.step(noise_pred, t, latents, variance_noise=zs, **extra_step_kwargs).prev_sample total_latents[b:b + batch_size] = latents return total_latents @torch.no_grad() def ddim_sample(self, x, cond, batch_size): timesteps = self.scheduler.timesteps for i, t in enumerate(tqdm(timesteps)): for b in range(0, x.shape[0], batch_size): x_batch = x[b:b + batch_size] model_input = x_batch cond_batch = cond.repeat(x_batch.shape[0], 1, 1) if self.sd_version == 'depth': depth_maps = torch.cat([self.depth_maps[b: b + batch_size]]) model_input = torch.cat([x_batch, depth_maps],dim=1) alpha_prod_t = self.scheduler.alphas_cumprod[t] alpha_prod_t_prev = ( self.scheduler.alphas_cumprod[timesteps[i + 1]] if i < len(timesteps) - 1 else self.scheduler.final_alpha_cumprod ) mu = alpha_prod_t ** 0.5 sigma = (1 - alpha_prod_t) ** 0.5 mu_prev = alpha_prod_t_prev ** 0.5 sigma_prev = (1 - alpha_prod_t_prev) ** 0.5 eps = self.unet(model_input, t, encoder_hidden_states=cond_batch).sample if self.sd_version != 'ControlNet' \ else self.controlnet_pred(x_batch, t, cond_batch, torch.cat([self.canny_cond[b: b + batch_size]])) pred_x0 = (x_batch - sigma * eps) / mu x[b:b + batch_size] = mu_prev * pred_x0 + sigma_prev * eps return x @torch.no_grad() def extract_latents(self, num_steps, save_path, batch_size, timesteps_to_save, inversion_prompt='', skip_steps=20, inversion_type='ddim', eta=1.0, reconstruction=False): self.scheduler.set_timesteps(num_steps) cond = self.get_text_embeds(inversion_prompt, "")[1].unsqueeze(0) latent_frames = self.latents if inversion_type == 'ddim': inverted_x= self.ddim_inversion(cond, latent_frames, save_path, batch_size=batch_size, save_latents=True if save_path else False, timesteps_to_save=timesteps_to_save) if reconstruction: latent_reconstruction = self.ddim_sample(inverted_x, cond, batch_size=batch_size) rgb_reconstruction = self.decode_latents(latent_reconstruction) return (self.frames, self.latents, self.total_inverted_latents), rgb_reconstruction else: return (self.frames, self.latents, self.total_inverted_latents), None elif inversion_type == 'ddpm': inverted_x, zs = self.ddpm_inversion(cond, latent_frames, save_path= save_path, batch_size=batch_size, save_latents=True if save_path else False, num_inversion_steps=num_steps, eta=eta, skip_steps=skip_steps) cond = self.encode_text(inversion_prompt) if reconstruction: latent_reconstruction = self.ddpm_sample(init_latents=inverted_x, cond=cond, batch_size=batch_size, num_inversion_steps=num_steps, skip_steps=skip_steps, eta=eta, zs_all=zs) rgb_reconstruction = self.decode_latents(latent_reconstruction) return (self.frames, self.latents, self.total_inverted_latents, self.noise_total), rgb_reconstruction else: return (self.frames, self.latents, self.total_inverted_latents, self.noise_total), None else: raise NotImplementedError() def compute_noise(scheduler, prev_latents, latents, timestep, noise_pred, eta): # 1. get previous step value (=t-1) prev_timestep = timestep - scheduler.config.num_train_timesteps // scheduler.num_inference_steps # 2. compute alphas, betas alpha_prod_t = scheduler.alphas_cumprod[timestep] alpha_prod_t_prev = ( scheduler.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else scheduler.final_alpha_cumprod ) beta_prod_t = 1 - alpha_prod_t # 3. compute predicted original sample from predicted noise also called # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf pred_original_sample = (latents - beta_prod_t ** (0.5) * noise_pred) / alpha_prod_t ** (0.5) # 4. Clip "predicted x_0" if scheduler.config.clip_sample: pred_original_sample = torch.clamp(pred_original_sample, -1, 1) # 5. compute variance: "sigma_t(η)" -> see formula (16) # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1) variance = scheduler._get_variance(timestep, prev_timestep) std_dev_t = eta * variance ** (0.5) # 6. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t ** 2) ** (0.5) * noise_pred # modifed so that updated xtm1 is returned as well (to avoid error accumulation) mu_xt = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction noise = (prev_latents - mu_xt) / (variance ** (0.5) * eta) return noise, mu_xt + (eta * variance ** 0.5) * noise def prep(opt): # timesteps to save if opt["sd_version"] == '2.1': model_key = "stabilityai/stable-diffusion-2-1-base" elif opt["sd_version"] == '2.0': model_key = "stabilityai/stable-diffusion-2-base" elif opt["sd_version"] == '1.5' or opt["sd_version"] == 'ControlNet': model_key = "runwayml/stable-diffusion-v1-5" elif opt["sd_version"] == 'depth': model_key = "stabilityai/stable-diffusion-2-depth" toy_scheduler = DDIMScheduler.from_pretrained(model_key, subfolder="scheduler") toy_scheduler.set_timesteps(opt["save_steps"]) timesteps_to_save, num_inference_steps = get_timesteps(toy_scheduler, num_inference_steps=opt["save_steps"], strength=1.0, device=device) seed_everything(opt["seed"]) if not opt["frames"]: # original non demo setting save_path = os.path.join(opt["save_dir"], f'inversion_{opt[inversion]}', f'sd_{opt["sd_version"]}', Path(opt["data_path"]).stem, f'steps_{opt["steps"]}', f'nframes_{opt["n_frames"]}') os.makedirs(os.path.join(save_path, f'latents'), exist_ok=True) if opt[inversion] == 'ddpm': os.makedirs(os.path.join(save_path, f'latents'), exist_ok=True) add_dict_to_yaml_file(os.path.join(opt["save_dir"], 'inversion_prompts.yaml'), Path(opt["data_path"]).stem, opt["inversion_prompt"]) # save inversion prompt in a txt file with open(os.path.join(save_path, 'inversion_prompt.txt'), 'w') as f: f.write(opt["inversion_prompt"]) else: save_path = None model = Preprocess(device, config, vae=vae, text_encoder=text_encoder, scheduler=scheduler, tokenizer=tokenizer, unet=unet) frames_and_latents, rgb_reconstruction = model.extract_latents( num_steps=model.config["steps"], save_path=save_path, batch_size=model.config["batch_size"], timesteps_to_save=timesteps_to_save, inversion_prompt=model.config["inversion_prompt"], inversion_type=model.config[inversion], skip_steps=model.config[skip_steps], reconstruction=model.config[reconstruct] ) if model.config[inversion] == 'ddpm': frames, latents, total_inverted_latents, zs = frames_and_latents return frames, latents, total_inverted_latents, zs, rgb_reconstruction else: frames, latents, total_inverted_latents = frames_and_latents return frames, latents, total_inverted_latents, rgb_reconstruction