Create generate.py

generate.py

@@ -0,0 +1,375 @@
+import torch.nn as nn
+import torch
+import numpy as np
+from tqdm import tqdm
+import os
+from transformers import logging
+
+from utils import CONTROLNET_DICT
+from utils import load_config, save_config
+from utils import get_controlnet_kwargs, get_frame_ids, get_latents_dir, init_model, seed_everything
+from utils import prepare_control, load_latent, load_video, prepare_depth, save_video
+from utils import register_time, register_attention_control, register_conv_control
+
+import vidtome
+
+# suppress partial model loading warning
+logging.set_verbosity_error()
+
+
+class Generator(nn.Module):
+    def __init__(self, pipe, scheduler, config):
+        super().__init__()
+
+        self.device = config.device
+        self.seed = config.seed
+
+
+
+        
+        self.model_key = config.model_key
+
+        self.config = config
+        gene_config = config.generation
+        float_precision = gene_config.float_precision if "float_precision" in gene_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.n_timesteps = gene_config.n_timesteps
+        scheduler.set_timesteps(gene_config.n_timesteps, device=self.device)
+        self.scheduler = scheduler
+
+        self.batch_size = 2
+        self.control = gene_config.control
+        self.use_depth = config.sd_version == "depth"
+        self.use_controlnet = self.control in CONTROLNET_DICT.keys()
+        self.use_pnp = self.control == "pnp"
+        if self.use_controlnet:
+            self.controlnet = pipe.controlnet
+            self.controlnet_scale = gene_config.control_scale
+        elif self.use_pnp:
+            pnp_f_t = int(gene_config.n_timesteps * gene_config.pnp_f_t)
+            pnp_attn_t = int(gene_config.n_timesteps * gene_config.pnp_attn_t)
+            self.batch_size += 1
+            self.init_pnp(conv_injection_t=pnp_f_t, qk_injection_t=pnp_attn_t)
+
+        self.chunk_size = gene_config.chunk_size
+        self.chunk_ord = gene_config.chunk_ord
+        self.merge_global = gene_config.merge_global
+        self.local_merge_ratio = gene_config.local_merge_ratio
+        self.global_merge_ratio = gene_config.global_merge_ratio
+        self.global_rand = gene_config.global_rand
+        self.align_batch = gene_config.align_batch
+
+        self.prompt = gene_config.prompt
+        self.negative_prompt = gene_config.negative_prompt
+        self.guidance_scale = gene_config.guidance_scale
+        self.save_frame = gene_config.save_frame
+
+        self.frame_height, self.frame_width = config.height, config.width
+        self.work_dir = config.work_dir
+
+        self.chunk_ord = gene_config.chunk_ord
+        if "mix" in self.chunk_ord:
+            self.perm_div = float(self.chunk_ord.split("-")[-1]) if "-" in self.chunk_ord else 3.
+            self.chunk_ord = "mix"
+        # Patch VidToMe to model
+        self.activate_vidtome()
+
+        if gene_config.use_lora:
+            self.pipe.load_lora_weights(**gene_config.lora)
+    
+    def activate_vidtome(self):
+        vidtome.apply_patch(self.pipe, self.local_merge_ratio, self.merge_global, self.global_merge_ratio, 
+            seed = self.seed, batch_size = self.batch_size, align_batch = self.use_pnp or self.align_batch, global_rand = self.global_rand)        
+
+    @torch.no_grad()
+    def get_text_embeds_input(self, prompt, negative_prompt):
+        text_embeds = self.get_text_embeds(
+            prompt, negative_prompt, self.device)
+        if self.use_pnp:
+            pnp_guidance_embeds = self.get_text_embeds("", device=self.device)
+            text_embeds = torch.cat(
+                [pnp_guidance_embeds, text_embeds], dim=0)
+        return text_embeds
+
+    @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 prepare_data(self, data_path, latent_path, frame_ids):
+        self.frames = load_video(data_path, self.frame_height,
+                                 self.frame_width, frame_ids=frame_ids, device=self.device)
+        self.init_noise = load_latent(
+            latent_path, t=self.scheduler.timesteps[0], frame_ids=frame_ids).to(self.dtype).to(self.device)
+
+        if self.use_depth:
+            self.depths = prepare_depth(
+                self.pipe, self.frames, frame_ids, self.work_dir).to(self.init_noise)
+
+        if self.use_controlnet:
+            self.controlnet_images = prepare_control(
+                self.control, self.frames, frame_ids, self.work_dir).to(self.init_noise)
+
+    @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
+    
+    def get_chunks(self, flen):
+        x_index = torch.arange(flen)
+
+        # The first chunk has a random length
+        rand_first = np.random.randint(0, self.chunk_size) + 1
+        chunks = x_index[rand_first:].split(self.chunk_size, dim=0)
+        chunks = [x_index[:rand_first]] + list(chunks) if len(chunks[0]) > 0 else [x_index[:rand_first]]
+        if np.random.rand() > 0.5:
+            chunks = chunks[::-1]
+        
+        # Chunk order only matter when we do global token merging
+        if self.merge_global == False:
+            return chunks
+
+        # Chunk order. "seq": sequential order. "rand": full permutation. "mix": partial permutation.
+        if self.chunk_ord == "rand":
+            order = torch.randperm(len(chunks))
+        elif self.chunk_ord == "mix":
+            randord = torch.randperm(len(chunks)).tolist()
+            rand_len = int(len(randord) / self.perm_div)
+            seqord = sorted(randord[rand_len:])
+            if rand_len > 0:
+                randord = randord[:rand_len]
+                if abs(seqord[-1] - randord[-1]) < abs(seqord[0] - randord[-1]):
+                    seqord = seqord[::-1]
+                order = randord + seqord
+            else:
+                order = seqord
+        else:
+            order = torch.arange(len(chunks))
+        chunks = [chunks[i] for i in order]
+        return chunks
+
+    @torch.no_grad()
+    def ddim_sample(self, x, conds):
+        print("[INFO] denoising frames...")
+        timesteps = self.scheduler.timesteps
+        noises = torch.zeros_like(x)
+
+        for i, t in enumerate(tqdm(timesteps, desc="Sampling")):
+            self.pre_iter(x, t)
+
+            # Split video into chunks and denoise
+            chunks = self.get_chunks(len(x))
+            for chunk in chunks:
+                torch.cuda.empty_cache()
+                noises[chunk] = self.pred_noise(
+                    x[chunk], conds, t, batch_idx=chunk)
+
+            x = self.pred_next_x(x, noises, t, i, inversion=False)
+
+            self.post_iter(x, t)
+        return x
+
+    def pre_iter(self, x, t):
+        if self.use_pnp:
+            # Prepare PnP
+            register_time(self, t.item())
+            cur_latents = load_latent(self.latent_path, t=t, frame_ids = self.frame_ids)
+            self.cur_latents = cur_latents
+
+    def post_iter(self, x, t):
+        if self.merge_global:
+            # Reset global tokens
+            vidtome.update_patch(self.pipe, global_tokens = None)
+
+    @torch.no_grad()
+    def pred_noise(self, x, cond, t, batch_idx=None):
+
+        flen = len(x)
+        text_embed_input = cond.repeat_interleave(flen, dim=0)
+
+        # For classifier-free guidance
+        latent_model_input = torch.cat([x, x])
+        batch_size = 2
+
+        if self.use_pnp:
+            # Cat latents from inverted source frames for PnP operation
+            source_latents = self.cur_latents
+            if batch_idx is not None:
+                source_latents = source_latents[batch_idx]
+            latent_model_input = torch.cat([source_latents.to(x), latent_model_input])
+            batch_size += 1
+
+        # For sd-depth model
+        if self.use_depth:
+            depth = self.depths
+            if batch_idx is not None:
+                depth = depth[batch_idx]
+            depth = depth.repeat(batch_size, 1, 1, 1)
+            latent_model_input = torch.cat([latent_model_input, depth.to(x)], dim=1)
+        
+        kwargs = dict()
+        # Compute controlnet outputs
+        if self.use_controlnet:
+            controlnet_cond = self.controlnet_images
+            if batch_idx is not None:
+                controlnet_cond = controlnet_cond[batch_idx]
+            controlnet_cond = controlnet_cond.repeat(batch_size, 1, 1, 1)
+            controlnet_kwargs = get_controlnet_kwargs(
+                self.controlnet, latent_model_input, text_embed_input, t, controlnet_cond, self.controlnet_scale)
+            kwargs.update(controlnet_kwargs)
+        # Pred noise!
+        eps = self.unet(latent_model_input, t, encoder_hidden_states=text_embed_input, **kwargs).sample
+        noise_pred_uncond, noise_pred_cond = eps.chunk(batch_size)[-2:]
+        # CFG
+        noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_cond - noise_pred_uncond)
+        return noise_pred
+
+    @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
+
+    def init_pnp(self, conv_injection_t, qk_injection_t):
+        qk_injection_timesteps = self.scheduler.timesteps[:qk_injection_t] if qk_injection_t >= 0 else []
+        conv_injection_timesteps = self.scheduler.timesteps[:conv_injection_t] if conv_injection_t >= 0 else []
+        register_attention_control(
+            self, qk_injection_timesteps, num_inputs=self.batch_size)
+        register_conv_control(
+            self, conv_injection_timesteps, num_inputs=self.batch_size)
+
+    def check_latent_exists(self, latent_path):
+        if self.use_pnp:
+            timesteps = self.scheduler.timesteps
+        else:
+            timesteps = [self.scheduler.timesteps[0]]
+
+        for ts in timesteps:
+            cur_latent_path = os.path.join(
+                latent_path, f'noisy_latents_{ts}.pt')
+            if not os.path.exists(cur_latent_path):
+                return False
+        return True
+
+    @torch.no_grad()
+    def __call__(self, data_path, latent_path, output_path, frame_ids):
+        self.scheduler.set_timesteps(self.n_timesteps)
+        latent_path = get_latents_dir(latent_path, self.model_key)
+        assert self.check_latent_exists(
+            latent_path), f"Required latent not found at {latent_path}. \
+                    Note: If using PnP as control, you need inversion latents saved \
+                     at each generation timestep."
+        
+        self.data_path = data_path
+        self.latent_path = latent_path
+        self.frame_ids = frame_ids
+        self.prepare_data(data_path, latent_path, frame_ids)
+
+        print(f"[INFO] initial noise latent shape: {self.init_noise.shape}")
+
+        for edit_name, edit_prompt in self.prompt.items():
+            print(f"[INFO] current prompt: {edit_prompt}")
+            conds = self.get_text_embeds_input(edit_prompt, self.negative_prompt)
+            # Comment this if you have enough GPU memory
+            clean_latent = self.ddim_sample(self.init_noise, conds)
+            torch.cuda.empty_cache()
+            clean_frames = self.decode_latents_batch(clean_latent)
+            cur_output_path = os.path.join(output_path, edit_name)
+            save_config(self.config, cur_output_path, gene = True)
+            save_video(clean_frames, cur_output_path, save_frame = self.save_frame)
+
+
+        
+
+
+if __name__ == "__main__":
+    config = load_config()
+    pipe, scheduler, model_key = init_model(
+        config.device, config.sd_version, config.model_key, config.generation.control, config.float_precision)
+    config.model_key = model_key
+    seed_everything(config.seed)
+    generator = Generator(pipe, scheduler, config)
+    frame_ids = get_frame_ids(
+        config.generation.frame_range, config.generation.frame_ids)
+    generator(config.input_path, config.generation.latents_path,
+              config.generation.output_path, frame_ids=frame_ids)