Update tokenflow_pnp.py

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: