Create invert.py

invert.py

@@ -0,0 +1,289 @@
+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)