import torch.nn as nn import torch from tqdm import tqdm import os from transformers import logging from utils import load_config, save_config from utils import get_controlnet_kwargs, get_latents_dir, init_model, seed_everything from utils import load_video, prepare_depth, save_frames, control_preprocess # suppress partial model loading warning logging.set_verbosity_error() class Inverter(nn.Module): def __init__(self, pipe, scheduler, config): super().__init__() self.device = config.device self.use_depth = config.sd_version == "depth" self.model_key = config.model_key self.config = config inv_config = config.inversion float_precision = inv_config.float_precision if "float_precision" in inv_config else config.float_precision if float_precision == "fp16": self.dtype = torch.float16 print("[INFO] float precision fp16. Use torch.float16.") else: self.dtype = torch.float32 print("[INFO] float precision fp32. Use torch.float32.") self.pipe = pipe self.vae = pipe.vae self.tokenizer = pipe.tokenizer self.unet = pipe.unet self.text_encoder = pipe.text_encoder if config.enable_xformers_memory_efficient_attention: try: pipe.enable_xformers_memory_efficient_attention() except ModuleNotFoundError: print("[WARNING] xformers not found. Disable xformers attention.") self.control = inv_config.control if self.control != "none": self.controlnet = pipe.controlnet self.controlnet_scale = inv_config.control_scale scheduler.set_timesteps(inv_config.save_steps) self.timesteps_to_save = scheduler.timesteps scheduler.set_timesteps(inv_config.steps) self.scheduler = scheduler self.prompt=inv_config.prompt self.recon=inv_config.recon self.save_latents=inv_config.save_intermediate self.use_blip=inv_config.use_blip self.steps=inv_config.steps self.batch_size = inv_config.batch_size self.force = inv_config.force self.n_frames = inv_config.n_frames self.frame_height, self.frame_width = config.height, config.width self.work_dir = config.work_dir @torch.no_grad() def get_text_embeds(self, prompt, negative_prompt=None, device="cuda"): text_input = self.tokenizer(prompt, padding='max_length', max_length=self.tokenizer.model_max_length, truncation=True, return_tensors='pt') text_embeddings = self.text_encoder(text_input.input_ids.to(device))[0] if negative_prompt is not None: uncond_input = self.tokenizer(negative_prompt, padding='max_length', max_length=self.tokenizer.model_max_length, return_tensors='pt') uncond_embeddings = self.text_encoder( uncond_input.input_ids.to(device))[0] text_embeddings = torch.cat([uncond_embeddings, text_embeddings]) return text_embeddings @torch.no_grad() def decode_latents(self, latents): with torch.autocast(device_type=self.device, dtype=self.dtype): latents = 1 / 0.18215 * latents imgs = self.vae.decode(latents).sample imgs = (imgs / 2 + 0.5).clamp(0, 1) return imgs @torch.no_grad() def decode_latents_batch(self, latents): imgs = [] batch_latents = latents.split(self.batch_size, dim = 0) for latent in batch_latents: imgs += [self.decode_latents(latent)] imgs = torch.cat(imgs) return imgs @torch.no_grad() def encode_imgs(self, imgs): with torch.autocast(device_type=self.device, dtype=self.dtype): imgs = 2 * imgs - 1 posterior = self.vae.encode(imgs).latent_dist latents = posterior.mean * 0.18215 return latents @torch.no_grad() def encode_imgs_batch(self, imgs): latents = [] batch_imgs = imgs.split(self.batch_size, dim = 0) for img in batch_imgs: latents += [self.encode_imgs(img)] latents = torch.cat(latents) return latents @torch.no_grad() def ddim_inversion(self, x, conds, save_path): print("[INFO] start DDIM Inversion!") timesteps = reversed(self.scheduler.timesteps) with torch.autocast(device_type=self.device, dtype=self.dtype): for i, t in enumerate(tqdm(timesteps)): noises = [] x_index = torch.arange(len(x)) batches = x_index.split(self.batch_size, dim = 0) for batch in batches: noise = self.pred_noise( x[batch], conds[batch], timesteps[i], batch_idx=batch) noises += [noise] noises = torch.cat(noises) x = self.pred_next_x(x, noises, t, i, inversion=True) if self.save_latents and t in self.timesteps_to_save: torch.save(x, os.path.join( save_path, f'noisy_latents_{t}.pt')) # Save inverted noise latents pth = os.path.join(save_path, f'noisy_latents_{t}.pt') torch.save(x, pth) print(f"[INFO] inverted latent saved to: {pth}") return x @torch.no_grad() def ddim_sample(self, x, conds): print("[INFO] reconstructing frames...") timesteps = self.scheduler.timesteps with torch.autocast(device_type=self.device, dtype=self.dtype): for i, t in enumerate(tqdm(timesteps)): noises = [] x_index = torch.arange(len(x)) batches = x_index.split(self.batch_size, dim = 0) for batch in batches: noise = self.pred_noise( x[batch], conds[batch], t, batch_idx=batch) noises += [noise] noises = torch.cat(noises) x = self.pred_next_x(x, noises, t, i, inversion=False) return x @torch.no_grad() def pred_noise(self, x, cond, t, batch_idx=None): # For sd-depth model if self.use_depth: depth = self.depths if batch_idx is not None: depth = depth[batch_idx] x = torch.cat([x, depth.to(x)], dim=1) kwargs = dict() # Compute controlnet outputs if self.control != "none": if batch_idx is None: controlnet_cond = self.controlnet_images else: controlnet_cond = self.controlnet_images[batch_idx] controlnet_kwargs = get_controlnet_kwargs(self.controlnet, x, cond, t, controlnet_cond, self.controlnet_scale) kwargs.update(controlnet_kwargs) eps = self.unet(x, t, encoder_hidden_states=cond, **kwargs).sample return eps @torch.no_grad() def pred_next_x(self, x, eps, t, i, inversion=False): if inversion: timesteps = reversed(self.scheduler.timesteps) else: timesteps = self.scheduler.timesteps alpha_prod_t = self.scheduler.alphas_cumprod[t] if inversion: alpha_prod_t_prev = ( self.scheduler.alphas_cumprod[timesteps[i - 1]] if i > 0 else self.scheduler.final_alpha_cumprod ) else: 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 if inversion: pred_x0 = (x - sigma_prev * eps) / mu_prev x = mu * pred_x0 + sigma * eps else: pred_x0 = (x - sigma * eps) / mu x = mu_prev * pred_x0 + sigma_prev * eps return x @torch.no_grad() def prepare_cond(self, prompts, n_frames): if isinstance(prompts, str): prompts = [prompts] * n_frames cond = self.get_text_embeds(prompts[0]) conds = torch.cat([cond] * n_frames) elif isinstance(prompts, list): cond_ls = [] for prompt in prompts: cond = self.get_text_embeds(prompt) cond_ls += [cond] conds = torch.cat(cond_ls) return conds, prompts def check_latent_exists(self, save_path): save_timesteps = [self.scheduler.timesteps[0]] if self.save_latents: save_timesteps += self.timesteps_to_save for ts in save_timesteps: latent_path = os.path.join( save_path, f'noisy_latents_{ts}.pt') if not os.path.exists(latent_path): return False return True @torch.no_grad() def __call__(self, data_path, save_path): self.scheduler.set_timesteps(self.steps) save_path = get_latents_dir(save_path, self.model_key) os.makedirs(save_path, exist_ok = True) if self.check_latent_exists(save_path) and not self.force: print(f"[INFO] inverted latents exist at: {save_path}. Skip inversion! Set 'inversion.force: True' to invert again.") return frames = load_video(data_path, self.frame_height, self.frame_width, device = self.device) frame_ids = list(range(len(frames))) if self.n_frames is not None: frame_ids = frame_ids[:self.n_frames] frames = frames[frame_ids] if self.use_depth: self.depths = prepare_depth(self.pipe, frames, frame_ids, self.work_dir) conds, prompts = self.prepare_cond(self.prompt, len(frames)) with open(os.path.join(save_path, 'inversion_prompts.txt'), 'w') as f: f.write('\n'.join(prompts)) if self.control != "none": images = control_preprocess( frames, self.control) self.controlnet_images = images.to(self.device) latents = self.encode_imgs_batch(frames) torch.cuda.empty_cache() print(f"[INFO] clean latents shape: {latents.shape}") inverted_x = self.ddim_inversion(latents, conds, save_path) save_config(self.config, save_path, inv = True) if self.recon: latent_reconstruction = self.ddim_sample(inverted_x, conds) torch.cuda.empty_cache() recon_frames = self.decode_latents_batch( latent_reconstruction) recon_save_path = os.path.join(save_path, 'recon_frames') save_frames(recon_frames, recon_save_path, frame_ids = frame_ids) if __name__ == "__main__": config = load_config() pipe, scheduler, model_key = init_model( config.device, config.sd_version, config.model_key, config.inversion.control, config.float_precision) config.model_key = model_key seed_everything(config.seed) inversion = Inverter(pipe, scheduler, config) inversion(config.input_path, config.inversion.save_path)