add ddpm inversion (#4)

- add ddpm inversion (e24d40d1c2c8663f66146b7451847baa9ba55c08)
- Update preprocess_utils.py (248a53d9dad0242311a3673b4b6c6e9e4188e6fd)
- Update tokenflow_pnp.py (2bd26712f2ff49f7fd34429b0a6d2677fa918139)

app.py

@@ -61,20 +61,22 @@ def prep(config):
         model_key = "stabilityai/stable-diffusion-2-depth"
     toy_scheduler = DDIMScheduler.from_pretrained(model_key, subfolder="scheduler")
     toy_scheduler.set_timesteps(config["save_steps"])
-    print("config[save_steps]", config["save_steps"])
     timesteps_to_save, num_inference_steps = get_timesteps(toy_scheduler, num_inference_steps=config["save_steps"],
                                                            strength=1.0,
                                                            device=device)
-    print("YOOOO timesteps to save", timesteps_to_save)
+
 
     # seed_everything(config["seed"])
     if not config["frames"]: # original non demo setting
         save_path = os.path.join(config["save_dir"],
+                                 f'inversion_{config[inversion]}',
                                  f'sd_{config["sd_version"]}',
                                  Path(config["data_path"]).stem,
                                  f'steps_{config["steps"]}',
                                  f'nframes_{config["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(config["save_dir"], 'inversion_prompts.yaml'), Path(config["data_path"]).stem, config["inversion_prompt"])    
         # save inversion prompt in a txt file
         with open(os.path.join(save_path, 'inversion_prompt.txt'), 'w') as f:
@@ -82,43 +84,53 @@ def prep(config):
     else:
         save_path = None
     
-    model = Preprocess(device, config,
+    model = Preprocess(device, 
+                       config,
                       vae=vae,
                       text_encoder=text_encoder,
                       scheduler=scheduler,
                       tokenizer=tokenizer,
                       unet=unet)
-    print(type(model.config["batch_size"]))
-    frames, latents, total_inverted_latents, rgb_reconstruction = model.extract_latents(
+
+    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"]
     )
 
-    
-    return frames, latents, total_inverted_latents, rgb_reconstruction
+    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
+
     
 def preprocess_and_invert(input_video,
                           frames,
                           latents,
                           inverted_latents,
+                          zs,
                           seed, 
                           randomize_seed,
                           do_inversion,
-                          # save_dir: str = "latents",
                           steps,
                           n_timesteps = 50,
                           batch_size: int = 8,
                           n_frames: int = 40,
                           inversion_prompt:str = '',
+                          skip_steps: int = 15,
                           
               ):
     sd_version = "2.1"
-    height = 512
+    height: int = 512
     weidth: int = 512
-    print("n timesteps", n_timesteps)
+
     if do_inversion or randomize_seed:
         preprocess_config = {}
         preprocess_config['H'] = height
@@ -134,30 +146,37 @@ def preprocess_and_invert(input_video,
         preprocess_config['frames'] = video_to_frames(input_video)
         preprocess_config['data_path'] = input_video.split(".")[0]
         
+        preprocess_config['inversion'] = 'ddpm'
+        preprocess_config['skip_steps'] = skip_steps
+        preprocess_config['reconstruct'] = False
+        
+        
 
         if randomize_seed:
             seed = randomize_seed_fn()
         seed_everything(seed)
         
-        frames, latents, total_inverted_latents, rgb_reconstruction = prep(preprocess_config)
-        print(total_inverted_latents.keys())
-        print(len(total_inverted_latents.keys()))
-        frames = gr.State(value=frames)
-        latents = gr.State(value=latents)
-        inverted_latents = gr.State(value=total_inverted_latents)
+        frames, latents, total_inverted_latents, zs, rgb_reconstruction = prep(preprocess_config)
+        
+        frames = gr.State(value = frames)
+        latents = gr.State(value = latents)
+        inverted_latents = gr.State(value = total_inverted_latents)
+        zs = gr.State(value = zs)
         do_inversion = False
    
-    return frames, latents, inverted_latents, do_inversion
+    return frames, latents, inverted_latents, zs, do_inversion
 
 
 def edit_with_pnp(input_video,
                   frames, 
                   latents,
                   inverted_latents,
+                  zs,
                   seed,
                   randomize_seed,
                   do_inversion,
                   steps,
+                  skip_steps: int = 15,
                   prompt: str = "a marble sculpture of a woman running, Venus de Milo",
                   # negative_prompt: str = "ugly, blurry, low res, unrealistic, unaesthetic",
                   pnp_attn_t: float = 0.5,
@@ -183,14 +202,18 @@ def edit_with_pnp(input_video,
     config["pnp_attn_t"] = pnp_attn_t
     config["pnp_f_t"] = pnp_f_t
     config["pnp_inversion_prompt"] = inversion_prompt
+    config["inversion"] = "ddpm"
+    config["skip_steps"] = skip_steps
+    
     
     
     if do_inversion:
-        frames, latents, inverted_latents, do_inversion =  preprocess_and_invert(
+        frames, latents, inverted_latents, zs, do_inversion =  preprocess_and_invert(
                           input_video,
                           frames,
                           latents,
                           inverted_latents,
+                          zs,
                           seed, 
                           randomize_seed,
                           do_inversion,
@@ -198,7 +221,8 @@ def edit_with_pnp(input_video,
                           n_timesteps,
                           batch_size,
                           n_frames,
-                          inversion_prompt)
+                          inversion_prompt,
+                          skip_steps)
         do_inversion = False
         
     
@@ -207,12 +231,13 @@ def edit_with_pnp(input_video,
     seed_everything(seed)
     
     
-    editor = TokenFlow(config=config,pipe=tokenflow_pipe, frames=frames.value, inverted_latents=inverted_latents.value)
+    editor = TokenFlow(config=config,pipe=tokenflow_pipe, frames=frames.value, inverted_latents=inverted_latents.value, zs= zs.value)
     edited_frames = editor.edit_video()
-
-    save_video(edited_frames, 'tokenflow_PnP_fps_30.mp4', fps=n_fps)
+    
+    edit_video_path = f'tokenflow_PnP_fps_{n_fps}.mp4'
+    save_video(edited_frames, edit_video_path, fps=n_fps)
     # path = export_to_video(edited_frames)
-    return 'tokenflow_PnP_fps_30.mp4', frames, latents, inverted_latents, do_inversion
+    return edit_video_path, frames, latents, inverted_latents, zs, do_inversion
 
 ########
 # demo #
@@ -238,6 +263,7 @@ with gr.Blocks(css="style.css") as demo:
     frames = gr.State()
     inverted_latents = gr.State()
     latents = gr.State()
+    zs = gr.State()
     do_inversion = gr.State(value=True)
 
     with gr.Row():
@@ -252,15 +278,7 @@ with gr.Blocks(css="style.css") as demo:
                             label="Describe your edited video",
                             max_lines=1, value=""
                         )
-    # with gr.Group(visible=False) as share_btn_container:
-        # with gr.Group(elem_id="share-btn-container"):
-        #     community_icon = gr.HTML(community_icon_html, visible=True)
-        #     loading_icon = gr.HTML(loading_icon_html, visible=False)
-        #     share_button = gr.Button("Share to community", elem_id="share-btn", visible=True)
-        
-   
-    # with gr.Row():
-    #     inversion_progress = gr.Textbox(visible=False, label="Inversion progress")
+    
                
     with gr.Row():
         run_button = gr.Button("Edit your video!", visible=True)
@@ -274,8 +292,10 @@ with gr.Blocks(css="style.css") as demo:
                         randomize_seed = gr.Checkbox(label='Randomize seed', value=False)
                         gudiance_scale = gr.Slider(label='Guidance Scale', minimum=1, maximum=30,
                                               value=7.5, step=0.5, interactive=True)
-                        steps = gr.Slider(label='Inversion steps', minimum=10, maximum=500,
-                                              value=200, step=1, interactive=True)
+                        steps = gr.Slider(label='Inversion steps', minimum=10, maximum=200,
+                                              value=50, step=1, interactive=True)
+                        skip_steps = gr.Slider(label='Skip Steps', minimum=5, maximum=25,
+                                              value=5, step=1, interactive=True)
                         
                     with gr.Column(min_width=100):
                         inversion_prompt = gr.Textbox(lines=1, label="Inversion prompt", interactive=True, placeholder="")
@@ -284,7 +304,7 @@ with gr.Blocks(css="style.css") as demo:
                         n_frames = gr.Slider(label='Num frames', minimum=2, maximum=200,
                                               value=24, step=1, interactive=True)
                         n_timesteps = gr.Slider(label='Diffusion steps', minimum=25, maximum=100,
-                                              value=25, step=25, interactive=True)
+                                              value=50, step=25, interactive=True)
                         n_fps = gr.Slider(label='Frames per second', minimum=1, maximum=60,
                                               value=10, step=1, interactive=True)
                         
@@ -300,6 +320,11 @@ with gr.Blocks(css="style.css") as demo:
         fn = reset_do_inversion,
         outputs = [do_inversion],
         queue = False)
+    
+    steps.change(
+        fn = reset_do_inversion,
+        outputs = [do_inversion],
+        queue = False)
 
     inversion_prompt.change(
         fn = reset_do_inversion,
@@ -326,6 +351,7 @@ with gr.Blocks(css="style.css") as demo:
                       frames,
                       latents,
                       inverted_latents,
+                      zs,
                       seed, 
                       randomize_seed,
                       do_inversion,
@@ -333,11 +359,13 @@ with gr.Blocks(css="style.css") as demo:
                       n_timesteps,
                       batch_size,
                       n_frames,
-                      inversion_prompt
+                      inversion_prompt,
+                      skip_steps
           ],
           outputs = [frames,
                      latents,
                      inverted_latents,
+                     zs,
                      do_inversion
               
           ])
@@ -347,10 +375,12 @@ with gr.Blocks(css="style.css") as demo:
                                frames, 
                               latents,
                               inverted_latents,
+                              zs,
                               seed,
                               randomize_seed,
                               do_inversion,
                               steps,
+                              skip_steps,
                               prompt,                             
                               pnp_attn_t,
                               pnp_f_t,
@@ -360,7 +390,7 @@ with gr.Blocks(css="style.css") as demo:
                               gudiance_scale,
                               inversion_prompt,
                               n_fps ],
-                                 outputs = [output_video, frames, latents, inverted_latents, do_inversion]
+                                 outputs = [output_video, frames, latents, inverted_latents, zs, do_inversion]
                                 )
 
     gr.Examples(
@@ -371,4 +401,5 @@ with gr.Blocks(css="style.css") as demo:
     )
 
 demo.queue()
+
 demo.launch()

preprocess_utils.py

@@ -1,5 +1,6 @@
 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()
 
@@ -12,6 +13,8 @@ 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):
@@ -64,7 +67,9 @@ class Preprocess(nn.Module):
         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)
@@ -163,14 +168,34 @@ class Preprocess(nn.Module):
         )[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_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]
-        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 = 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
 
@@ -192,7 +217,7 @@ class Preprocess(nn.Module):
         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 * 0.18215)
+            latents.append(latent * self.vae.config.scaling_factor)
         latents = torch.cat(latents)
         return latents
 
@@ -264,6 +289,137 @@ class Preprocess(nn.Module):
             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):
@@ -295,6 +451,8 @@ class Preprocess(nn.Module):
                 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, 
@@ -303,31 +461,89 @@ class Preprocess(nn.Module):
                         batch_size,
                         timesteps_to_save,
                         inversion_prompt='',
-                       reconstruct=False):
+                        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
-        print("latent_frames", latent_frames.shape)
-  
-        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 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()
 
-       
-        # print("total_inverted_latents", len(total_inverted_latents.keys()))
-        
-        if reconstruct:
-            latent_reconstruction = self.ddim_sample(inverted_x, cond, batch_size=batch_size)
+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
 
-            rgb_reconstruction = self.decode_latents(latent_reconstruction)
-            return self.frames, self.latents, self.total_inverted_latents, rgb_reconstruction
+    # 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)
 
-        return self.frames, self.latents, self.total_inverted_latents, None
+    # 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
@@ -348,11 +564,14 @@ def prep(opt):
     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:
@@ -360,16 +579,31 @@ def prep(opt):
     else:
         save_path = None
     
-    model = Preprocess(device, opt)
+    model = Preprocess(device, 
+                       config,
+                      vae=vae,
+                      text_encoder=text_encoder,
+                      scheduler=scheduler,
+                      tokenizer=tokenizer,
+                      unet=unet)
   
-    frames, latents, total_inverted_latents, rgb_reconstruction = model.extract_latents(
+    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]
     )
 
-    
-    return frames, latents, total_inverted_latents, rgb_reconstruction
+    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
+
+
 

tokenflow_pnp.py

@@ -9,6 +9,7 @@ import torchvision.transforms as T
 import argparse
 from PIL import Image
 import yaml
+import inspect
 from tqdm import tqdm
 from transformers import logging
 from diffusers import DDIMScheduler, StableDiffusionPipeline
@@ -25,9 +26,9 @@ VAE_BATCH_SIZE = 10
 class TokenFlow(nn.Module):
     def __init__(self, config, 
                  pipe,
-                 frames=None,
-                # latents = None,
-                inverted_latents = None):
+                 frames = None,
+                 inverted_latents = None, #X0,...,XT,
+                 zs = None):
         super().__init__()
         self.config = config
         self.device = config["device"]
@@ -61,7 +62,16 @@ class TokenFlow(nn.Module):
         print('SD model loaded')
         
         # data
-        self.frames, self.inverted_latents  = frames, inverted_latents
+        self.inversion = config['inversion']
+        if self.inversion == 'ddpm':
+            self.skip_steps = config['skip_steps']
+            self.eta = 1.0
+        else:
+            self.eta = 0.0
+        self.extra_step_kwargs = self.prepare_extra_step_kwargs(self.eta)
+        
+        # data
+        self.frames, self.inverted_latents, self.zs  = frames, inverted_latents, zs
         self.latents_path = self.get_latents_path()
         
         # load frames
@@ -120,15 +130,13 @@ class TokenFlow(nn.Module):
 
     def get_latents_path(self):
         read_from_files = self.frames is None
-        # read_from_files = True
         if read_from_files: 
             latents_path = os.path.join(self.config["latents_path"], f'sd_{self.config["sd_version"]}',
                                  Path(self.config["data_path"]).stem, f'steps_{self.config["n_inversion_steps"]}')
             latents_path = [x for x in glob.glob(f'{latents_path}/*') if '.' not in Path(x).name]
             n_frames = [int([x for x in latents_path[i].split('/') if 'nframes' in x][0].split('_')[1]) for i in range(len(latents_path))]
-            print("n_frames", n_frames)
             latents_path = latents_path[np.argmax(n_frames)]
-            print("latents_path", latents_path)
+
             self.config["n_frames"] = min(max(n_frames), self.config["n_frames"])
             
         else:
@@ -138,9 +146,8 @@ class TokenFlow(nn.Module):
         if self.config["n_frames"] % self.config["batch_size"] != 0:
             # make n_frames divisible by batch_size
             self.config["n_frames"] = self.config["n_frames"] - (self.config["n_frames"] % self.config["batch_size"])
-        print("Number of frames: ", self.config["n_frames"])
+
         if read_from_files:
-            print("YOOOOOOO", os.path.join(latents_path, 'latents'))
             return os.path.join(latents_path, 'latents')
         else:
             return None
@@ -206,37 +213,61 @@ class TokenFlow(nn.Module):
         # encode to latents
         latents = self.encode_imgs(frames, deterministic=True).to(torch.float16).to(self.device)
         # get noise
-        eps = self.get_ddim_eps(latents, range(self.config["n_frames"])).to(torch.float16).to(self.device)
+        if self.inversion == 'ddim':
+            eps = self.get_ddim_eps(latents, range(self.config["n_frames"])).to(torch.float16).to(self.device)
+        elif self.inversion == 'ddpm':
+            eps = self.get_ddpm_noise()
+        else:
+            raise NotImplementedError()
+        
         if not read_from_files:
             return None, frames, latents, eps
         return paths, frames, latents, eps
 
     def get_ddim_eps(self, latent, indices):
         read_from_files = self.inverted_latents is None
-        # read_from_files = True
         if read_from_files:
             noisest = max([int(x.split('_')[-1].split('.')[0]) for x in glob.glob(os.path.join(self.latents_path, f'noisy_latents_*.pt'))])
-            print("noisets:", noisest)
-            print("indecies:", indices)
             latents_path = os.path.join(self.latents_path, f'noisy_latents_{noisest}.pt')
             noisy_latent = torch.load(latents_path)[indices].to(self.device)
-            
-            # path = os.path.join('test_latents', f'noisy_latents_{noisest}.pt')
-            # f_noisy_latent = torch.load(path)[indices].to(self.device)
-            # print(f_noisy_latent==noisy_latent)
         else:
             noisest = max([int(key.split("_")[-1]) for key in self.inverted_latents.keys()])
-            print("noisets:", noisest)
-            print("indecies:", indices)
             noisy_latent = self.inverted_latents[f'noisy_latents_{noisest}'][indices]
 
         alpha_prod_T = self.scheduler.alphas_cumprod[noisest]
         mu_T, sigma_T = alpha_prod_T ** 0.5, (1 - alpha_prod_T) ** 0.5
         eps = (noisy_latent - mu_T * latent) / sigma_T
         return eps
+    
+    def get_ddpm_noise(self):
+        read_from_files = self.inverted_latents is None
+        idx_to_t = {int(k): int(v) for k, v in enumerate(self.scheduler.timesteps)}
+        t = idx_to_t[self.skip_steps]
+        if read_from_files:
+            x0_path = os.path.join(self.latents_path, f'noisy_latents_{t}.pt')
+            zs_path = os.path.join(self.latents_path, f'noise_total.pt')
+            x0 = torch.load(x0_path)[:self.config["n_frames"]].to(self.device)
+            zs = torch.load(zs_path)[self.skip_steps:, :self.config["n_frames"]].to(self.device)
+        else:
+            x0 = self.inverted_latents[f'noisy_latents_{t}'][:self.config["n_frames"]].to(self.device)
+            zs = self.zs[self.skip_steps:, :self.config["n_frames"]].to(self.device)
+        return x0, 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
+
+        return extra_step_kwargs
 
     @torch.no_grad()
-    def denoise_step(self, x, t, indices):
+    def denoise_step(self, x, t, indices, zs=None):
         # register the time step and features in pnp injection modules
         read_files = self.inverted_latents is None
 
@@ -264,21 +295,31 @@ class TokenFlow(nn.Module):
         noise_pred = noise_pred_uncond + self.config["guidance_scale"] * (noise_pred_cond - noise_pred_uncond)
 
         # compute the denoising step with the reference model
-        denoised_latent = self.scheduler.step(noise_pred, t, x)['prev_sample']
+        denoised_latent = self.scheduler.step(noise_pred, t, x, variance_noise=zs, **self.extra_step_kwargs)[
+            'prev_sample']
         return denoised_latent
     
     @torch.autocast(dtype=torch.float16, device_type='cuda')
-    def batched_denoise_step(self, x, t, indices):
+    def batched_denoise_step(self, x, t, indices, zs=None):
         batch_size = self.config["batch_size"]
         denoised_latents = []
-        pivotal_idx = torch.randint(batch_size, (len(x)//batch_size,)) + torch.arange(0,len(x),batch_size) 
-            
+        pivotal_idx = torch.randint(batch_size, (len(x) // batch_size,)) + torch.arange(0, len(x), batch_size)
+
         register_pivotal(self, True)
-        self.denoise_step(x[pivotal_idx], t, indices[pivotal_idx])
+        if zs is None:
+            zs_input = None
+        else:
+            zs_input = zs[pivotal_idx]
+        self.denoise_step(x[pivotal_idx], t, indices[pivotal_idx], zs_input)
         register_pivotal(self, False)
         for i, b in enumerate(range(0, len(x), batch_size)):
             register_batch_idx(self, i)
-            denoised_latents.append(self.denoise_step(x[b:b + batch_size], t, indices[b:b + batch_size]))
+            if zs is None:
+                zs_input = None
+            else:
+                zs_input = zs[b:b + batch_size]
+            denoised_latents.append(self.denoise_step(x[b:b + batch_size], t, indices[b:b + batch_size]
+                                                      , zs_input))
         denoised_latents = torch.cat(denoised_latents)
         return denoised_latents
 
@@ -309,7 +350,13 @@ class TokenFlow(nn.Module):
         
         self.init_method(conv_injection_t=pnp_f_t, qk_injection_t=pnp_attn_t)
         
-        noisy_latents = self.scheduler.add_noise(self.latents, self.eps, self.scheduler.timesteps[0])
+        if self.inversion == 'ddim':
+            noisy_latents = self.scheduler.add_noise(self.latents, self.eps, self.scheduler.timesteps[0])
+        elif self.inversion == 'ddpm':
+            noisy_latents = self.eps[0]
+        else:
+            raise NotImplementedError()
+            
         edited_frames = self.sample_loop(noisy_latents, torch.arange(self.config["n_frames"]))
         
         if save_files:
@@ -321,12 +368,24 @@ class TokenFlow(nn.Module):
             return edited_frames
 
     def sample_loop(self, x, indices):
-        save_files  = self.inverted_latents is None # if we're in the original non-demo setting
-        # save_files = True
+        save_files  = self.inverted_latents is None # if we're in the original non-demo settinge
         if save_files:
             os.makedirs(f'{self.config["output_path"]}/img_ode', exist_ok=True)
-        for i, t in enumerate(tqdm(self.scheduler.timesteps, desc="Sampling")):
-                x = self.batched_denoise_step(x, t, indices)
+        
+        timesteps = self.scheduler.timesteps
+        if self.inversion == 'ddpm':
+            zs_total = self.eps[1]
+
+            t_to_idx = {int(v): k for k, v in enumerate(timesteps[-zs_total.shape[0]:])}
+            timesteps = timesteps[-zs_total.shape[0]:]
+            
+        for i, t in enumerate(tqdm(timesteps, desc="Sampling")):
+            if self.inversion == 'ddpm':
+                idx = t_to_idx[int(t)]
+                zs = zs_total[idx]
+            else:
+                zs = None
+            x = self.batched_denoise_step(x, t, indices, zs)
         
         decoded_latents = self.decode_latents(x)
         if save_files: