├── .gitignore
├── LICENSE
├── README.md
├── assets
    └── videos.gif
├── configs
    ├── config_pnp.yaml
    └── config_sdedit.yaml
├── data
    ├── wolf.mp4
    └── woman-running.mp4
├── preprocess.py
├── requirements.txt
├── run_tokenflow_pnp.py
├── run_tokenflow_sdedit.py
├── tokenflow_utils.py
└── util.py


/.gitignore:
--------------------------------------------------------------------------------
1 | .DS_Store
2 | 


--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
 1 | MIT License
 2 | 
 3 | Copyright (c) 2023 Michal Geyer, Omer Bar-Tal, Shai Bagon, Tali Dekel
 4 | 
 5 | Permission is hereby granted, free of charge, to any person obtaining a copy
 6 | of this software and associated documentation files (the "Software"), to deal
 7 | in the Software without restriction, including without limitation the rights
 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 | 
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 | 
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
22 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # TokenFlow: Consistent Diffusion Features for Consistent Video Editing (ICLR 2024)
 2 | ## [<a href="https://diffusion-tokenflow.github.io/" target="_blank">Project Page</a>]
 3 | 
 4 | [![arXiv](https://img.shields.io/badge/arXiv-TokenFlow-b31b1b.svg)](https://arxiv.org/abs/2307.10373) [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/weizmannscience/tokenflow)
 5 | ![Pytorch](https://img.shields.io/badge/PyTorch->=1.10.0-Red?logo=pytorch)
 6 | 
 7 | 
 8 | 
 9 | [//]: # ([![Replicate]&#40;https://replicate.com/cjwbw/multidiffusion/badge&#41;]&#40;https://replicate.com/cjwbw/multidiffusion&#41;)
10 | 
11 | [//]: # ([![Hugging Face Spaces]&#40;https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue&#41;]&#40;https://huggingface.co/spaces/weizmannscience/text2live&#41;)
12 | 
13 | 
14 | 
15 | 
16 | https://github.com/omerbt/TokenFlow/assets/52277000/93dccd63-7e9a-4540-a941-31962361b0bb
17 | 
18 | 
19 | **TokenFlow** is a framework that enables consistent video editing, using a pre-trained text-to-image diffusion model, without any further training or finetuning.
20 | 
21 | [//]: # (as described in <a href="https://arxiv.org/abs/2302.08113" target="_blank">&#40;link to paper&#41;</a>.)
22 | 
23 | [//]: # (. It can be used for localized and global edits that change the texture of existing objects or augment the scene with semi-transparent effects &#40;e.g. smoke, fire, snow&#41;.)
24 | 
25 | [//]: # (### Abstract)
26 | >The generative AI revolution has been recently expanded to videos. Nevertheless, current state-of-the-art video mod- els are still lagging behind image models in terms of visual quality and user control over the generated content. In this work, we present a framework that harnesses the power of a text-to-image diffusion model for the task of text-driven video editing. Specifically, given a source video and a target text-prompt, our method generates a high-quality video that adheres to the target text, while preserving the spatial lay- out and dynamics of the input video. Our method is based on our key observation that consistency in the edited video can be obtained by enforcing consistency in the diffusion feature space. We achieve this by explicitly propagating diffusion features based on inter-frame correspondences, readily available in the model. Thus, our framework does not require any training or fine-tuning, and can work in con- junction with any off-the-shelf text-to-image editing method. We demonstrate state-of-the-art editing results on a variety of real-world videos.
27 | 
28 | For more see the [project webpage](https://diffusion-tokenflow.github.io).
29 | 
30 | ## Sample results
31 | 
32 | <td><img src="assets/videos.gif"></td>
33 | 
34 | ## Environment
35 | ```
36 | conda create -n tokenflow python=3.9
37 | conda activate tokenflow
38 | pip install -r requirements.txt
39 | ```
40 | ## Preprocess
41 | 
42 | Preprocess you video by running using the following command:
43 | ```
44 | python preprocess.py --data_path <data/myvideo.mp4> \
45 |                      --inversion_prompt <'' or a string describing the video content>
46 | ```
47 | Additional arguments:
48 | ```
49 |                      --save_dir <latents>
50 |                      --H <video height>
51 |                      --W <video width>
52 |                      --sd_version <Stable-Diffusion version>
53 |                      --steps <number of inversion steps>
54 |                      --save_steps <number of sampling steps that will be used later for editing>
55 |                      --n_frames <number of frames>
56 |                      
57 | ```
58 | more information on the arguments can be found here.
59 | 
60 | ### Note: 
61 | The video reconstruction will be saved as inverted.mp4. A good reconstruction is required for successfull editing with our method.
62 | 
63 | ## Editing
64 | 
65 | - TokenFlow is designed for structure-preserving edits. 
66 | - Our method is built on top of an image editing technique (e.g., Plug-and-Play, ControlNet, etc.) - therefore, it is important to ensure that the edit works with the chosen base technique. 
67 | - The LDM decoder may introduce some jitterness, depending on the original video. 
68 | 
69 | To edit your video, first create a yaml config as in ``configs/config_pnp.yaml``.
70 | Then run 
71 | ```
72 | python run_tokenflow_pnp.py
73 | ```
74 | 
75 | Similarly, if you want to use ControlNet or SDEedit, create a yaml config as in ``config/config_controlnet.yaml`` or ```configs/config_SDEdit.yaml``` and run ```python run_tokenflow_controlnet.py``` or ``python run_tokenflow_SDEdit.py`` respectivly.
76 | 
77 | 
78 | ## Citation
79 | ```
80 | @article{tokenflow2023,
81 |         title = {TokenFlow: Consistent Diffusion Features for Consistent Video Editing},
82 |         author = {Geyer, Michal and Bar-Tal, Omer and Bagon, Shai and Dekel, Tali},
83 |         journal={arXiv preprint arxiv:2307.10373},
84 |         year={2023}
85 |         }
86 | ```
87 | 
88 | 


--------------------------------------------------------------------------------
/assets/videos.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/omerbt/TokenFlow/5dd6a696277e9d18c23adc034dc4c052ccd64cb5/assets/videos.gif


--------------------------------------------------------------------------------
/configs/config_pnp.yaml:
--------------------------------------------------------------------------------
 1 | # general
 2 | seed: 1
 3 | device: 'cuda'
 4 | output_path: 'tokenflow-results'
 5 | 
 6 | # data
 7 | data_path: 'data/woman-running'
 8 | latents_path: 'latents' # should be the same as 'save_dir' arg used in preprocess
 9 | n_inversion_steps: 500 # for retrieving the latents of the inversion
10 | n_frames: 40
11 | 
12 | # diffusion
13 | sd_version: '2.1'
14 | guidance_scale: 7.5
15 | n_timesteps: 50
16 | prompt: "a marble sculpture of a woman running, Venus de Milo"
17 | negative_prompt: "ugly, blurry, low res, unrealistic, unaesthetic"
18 | batch_size: 8
19 | 
20 | # pnp params -- injection thresholds ∈ [0, 1]
21 | pnp_attn_t: 0.5
22 | pnp_f_t: 0.8
23 | 


--------------------------------------------------------------------------------
/configs/config_sdedit.yaml:
--------------------------------------------------------------------------------
 1 | # general
 2 | seed: 1
 3 | device: 'cuda'
 4 | output_path: 'tokenflow-results'
 5 | 
 6 | # data
 7 | data_path: 'wolf-part'
 8 | latents_path: 'latents' # should be the same as 'save_dir' arg used in preprocess
 9 | n_inversion_steps: 100 # for retrieving the latents of the inversion
10 | n_frames: 40
11 | 
12 | # diffusion
13 | sd_version: '2.1'
14 | guidance_scale: 7.5
15 | n_timesteps: 50
16 | prompt: a shiny silver robotic wolf
17 | negative_prompt: "ugly, blurry, low res, unrealistic, unaesthetic"
18 | batch_size: 8
19 | 
20 | # SDEdit params 
21 | start: 0.9 # start sampling from t = start * 1000
22 | use_ddim_noise: True # use ddim noise to noise the images. better structure preservation


--------------------------------------------------------------------------------
/data/wolf.mp4:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/omerbt/TokenFlow/5dd6a696277e9d18c23adc034dc4c052ccd64cb5/data/wolf.mp4


--------------------------------------------------------------------------------
/data/woman-running.mp4:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/omerbt/TokenFlow/5dd6a696277e9d18c23adc034dc4c052ccd64cb5/data/woman-running.mp4


--------------------------------------------------------------------------------
/preprocess.py:
--------------------------------------------------------------------------------
  1 | from transformers import CLIPTextModel, CLIPTokenizer, logging
  2 | from diffusers import AutoencoderKL, UNet2DConditionModel, DDIMScheduler
  3 | # suppress partial model loading warning
  4 | logging.set_verbosity_error()
  5 | 
  6 | import os
  7 | from tqdm import tqdm, trange
  8 | import torch
  9 | import torch.nn as nn
 10 | import argparse
 11 | from torchvision.io import write_video
 12 | from pathlib import Path
 13 | from util import *
 14 | import torchvision.transforms as T
 15 | 
 16 | 
 17 | def get_timesteps(scheduler, num_inference_steps, strength, device):
 18 |     # get the original timestep using init_timestep
 19 |     init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
 20 | 
 21 |     t_start = max(num_inference_steps - init_timestep, 0)
 22 |     timesteps = scheduler.timesteps[t_start:]
 23 | 
 24 |     return timesteps, num_inference_steps - t_start
 25 | 
 26 | 
 27 | class Preprocess(nn.Module):
 28 |     def __init__(self, device, opt, hf_key=None):
 29 |         super().__init__()
 30 | 
 31 |         self.device = device
 32 |         self.sd_version = opt.sd_version
 33 |         self.use_depth = False
 34 | 
 35 |         print(f'[INFO] loading stable diffusion...')
 36 |         if hf_key is not None:
 37 |             print(f'[INFO] using hugging face custom model key: {hf_key}')
 38 |             model_key = hf_key
 39 |         elif self.sd_version == '2.1':
 40 |             model_key = "stabilityai/stable-diffusion-2-1-base"
 41 |         elif self.sd_version == '2.0':
 42 |             model_key = "stabilityai/stable-diffusion-2-base"
 43 |         elif self.sd_version == '1.5' or self.sd_version == 'ControlNet':
 44 |             model_key = "runwayml/stable-diffusion-v1-5"
 45 |         elif self.sd_version == 'depth':
 46 |             model_key = "stabilityai/stable-diffusion-2-depth"
 47 |         else:
 48 |             raise ValueError(f'Stable-diffusion version {self.sd_version} not supported.')
 49 |         self.model_key = model_key
 50 |         # Create model
 51 |         self.vae = AutoencoderKL.from_pretrained(model_key, subfolder="vae", revision="fp16",
 52 |                                                  torch_dtype=torch.float16).to(self.device)
 53 |         self.tokenizer = CLIPTokenizer.from_pretrained(model_key, subfolder="tokenizer")
 54 |         self.text_encoder = CLIPTextModel.from_pretrained(model_key, subfolder="text_encoder", revision="fp16",
 55 |                                                           torch_dtype=torch.float16).to(self.device)
 56 |         self.unet = UNet2DConditionModel.from_pretrained(model_key, subfolder="unet", revision="fp16",
 57 |                                                    torch_dtype=torch.float16).to(self.device)
 58 |         self.paths, self.frames, self.latents = self.get_data(opt.data_path, opt.n_frames)
 59 |         
 60 |         if self.sd_version == 'ControlNet':
 61 |             from diffusers import ControlNetModel, StableDiffusionControlNetPipeline
 62 |             controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16).to(self.device)
 63 |             control_pipe = StableDiffusionControlNetPipeline.from_pretrained(
 64 |                 "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
 65 |             ).to(self.device)
 66 |             self.unet = control_pipe.unet
 67 |             self.controlnet = control_pipe.controlnet
 68 |             self.canny_cond = self.get_canny_cond()
 69 |         elif self.sd_version == 'depth':
 70 |             self.depth_maps = self.prepare_depth_maps()
 71 |         self.scheduler = DDIMScheduler.from_pretrained(model_key, subfolder="scheduler")
 72 |         
 73 |         # self.unet.enable_xformers_memory_efficient_attention()
 74 |         print(f'[INFO] loaded stable diffusion!')
 75 |         
 76 |     @torch.no_grad()   
 77 |     def prepare_depth_maps(self, model_type='DPT_Large', device='cuda'):
 78 |         depth_maps = []
 79 |         midas = torch.hub.load("intel-isl/MiDaS", model_type)
 80 |         midas.to(device)
 81 |         midas.eval()
 82 | 
 83 |         midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")
 84 | 
 85 |         if model_type == "DPT_Large" or model_type == "DPT_Hybrid":
 86 |             transform = midas_transforms.dpt_transform
 87 |         else:
 88 |             transform = midas_transforms.small_transform
 89 | 
 90 |         for i in range(len(self.paths)):
 91 |             img = cv2.imread(self.paths[i])
 92 |             img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
 93 | 
 94 |             latent_h = img.shape[0] // 8
 95 |             latent_w = img.shape[1] // 8
 96 |             
 97 |             input_batch = transform(img).to(device)
 98 |             prediction = midas(input_batch)
 99 | 
100 |             depth_map = torch.nn.functional.interpolate(
101 |                 prediction.unsqueeze(1),
102 |                 size=(latent_h, latent_w),
103 |                 mode="bicubic",
104 |                 align_corners=False,
105 |             )
106 |             depth_min = torch.amin(depth_map, dim=[1, 2, 3], keepdim=True)
107 |             depth_max = torch.amax(depth_map, dim=[1, 2, 3], keepdim=True)
108 |             depth_map = 2.0 * (depth_map - depth_min) / (depth_max - depth_min) - 1.0
109 |             depth_maps.append(depth_map)
110 | 
111 |         return torch.cat(depth_maps).to(self.device).to(torch.float16)
112 |     
113 |     @torch.no_grad()
114 |     def get_canny_cond(self):
115 |         canny_cond = []
116 |         for image in self.frames.cpu().permute(0, 2, 3, 1):
117 |             image = np.uint8(np.array(255 * image))
118 |             low_threshold = 100
119 |             high_threshold = 200
120 | 
121 |             image = cv2.Canny(image, low_threshold, high_threshold)
122 |             image = image[:, :, None]
123 |             image = np.concatenate([image, image, image], axis=2)
124 |             image = torch.from_numpy((image.astype(np.float32) / 255.0))
125 |             canny_cond.append(image)
126 |         canny_cond = torch.stack(canny_cond).permute(0, 3, 1, 2).to(self.device).to(torch.float16)
127 |         return canny_cond
128 |     
129 |     def controlnet_pred(self, latent_model_input, t, text_embed_input, controlnet_cond):
130 |         down_block_res_samples, mid_block_res_sample = self.controlnet(
131 |             latent_model_input,
132 |             t,
133 |             encoder_hidden_states=text_embed_input,
134 |             controlnet_cond=controlnet_cond,
135 |             conditioning_scale=1,
136 |             return_dict=False,
137 |         )
138 |         
139 |         # apply the denoising network
140 |         noise_pred = self.unet(
141 |             latent_model_input,
142 |             t,
143 |             encoder_hidden_states=text_embed_input,
144 |             cross_attention_kwargs={},
145 |             down_block_additional_residuals=down_block_res_samples,
146 |             mid_block_additional_residual=mid_block_res_sample,
147 |             return_dict=False,
148 |         )[0]
149 |         return noise_pred
150 |     
151 |     @torch.no_grad()
152 |     def get_text_embeds(self, prompt, negative_prompt, device="cuda"):
153 |         text_input = self.tokenizer(prompt, padding='max_length', max_length=self.tokenizer.model_max_length,
154 |                                     truncation=True, return_tensors='pt')
155 |         text_embeddings = self.text_encoder(text_input.input_ids.to(device))[0]
156 |         uncond_input = self.tokenizer(negative_prompt, padding='max_length', max_length=self.tokenizer.model_max_length,
157 |                                       return_tensors='pt')
158 |         uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(device))[0]
159 |         text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
160 |         return text_embeddings
161 | 
162 |     @torch.no_grad()
163 |     def decode_latents(self, latents):
164 |         decoded = []
165 |         batch_size = 8
166 |         for b in range(0, latents.shape[0], batch_size):
167 |                 latents_batch = 1 / 0.18215 * latents[b:b + batch_size]
168 |                 imgs = self.vae.decode(latents_batch).sample
169 |                 imgs = (imgs / 2 + 0.5).clamp(0, 1)
170 |                 decoded.append(imgs)
171 |         return torch.cat(decoded)
172 | 
173 |     @torch.no_grad()
174 |     def encode_imgs(self, imgs, batch_size=10, deterministic=True):
175 |         imgs = 2 * imgs - 1
176 |         latents = []
177 |         for i in range(0, len(imgs), batch_size):
178 |             posterior = self.vae.encode(imgs[i:i + batch_size]).latent_dist
179 |             latent = posterior.mean if deterministic else posterior.sample()
180 |             latents.append(latent * 0.18215)
181 |         latents = torch.cat(latents)
182 |         return latents
183 | 
184 |     def get_data(self, frames_path, n_frames):
185 |         # load frames
186 |         paths =  [f"{frames_path}/%05d.png" % i for i in range(n_frames)]
187 |         if not os.path.exists(paths[0]):
188 |             paths = [f"{frames_path}/%05d.jpg" % i for i in range(n_frames)]
189 |         self.paths = paths
190 |         frames = [Image.open(path).convert('RGB') for path in paths]
191 |         if frames[0].size[0] == frames[0].size[1]:
192 |             frames = [frame.resize((512, 512), resample=Image.Resampling.LANCZOS) for frame in frames]
193 |         frames = torch.stack([T.ToTensor()(frame) for frame in frames]).to(torch.float16).to(self.device)
194 |         # encode to latents
195 |         latents = self.encode_imgs(frames, deterministic=True).to(torch.float16).to(self.device)
196 |         return paths, frames, latents
197 | 
198 |     @torch.no_grad()
199 |     def ddim_inversion(self, cond, latent_frames, save_path, batch_size, save_latents=True, timesteps_to_save=None):
200 |         timesteps = reversed(self.scheduler.timesteps)
201 |         timesteps_to_save = timesteps_to_save if timesteps_to_save is not None else timesteps
202 |         for i, t in enumerate(tqdm(timesteps)):
203 |             for b in range(0, latent_frames.shape[0], batch_size):
204 |                 x_batch = latent_frames[b:b + batch_size]
205 |                 model_input = x_batch
206 |                 cond_batch = cond.repeat(x_batch.shape[0], 1, 1)
207 |                 if self.sd_version == 'depth':
208 |                     depth_maps = torch.cat([self.depth_maps[b: b + batch_size]])
209 |                     model_input = torch.cat([x_batch, depth_maps],dim=1)
210 |                                                                     
211 |                 alpha_prod_t = self.scheduler.alphas_cumprod[t]
212 |                 alpha_prod_t_prev = (
213 |                     self.scheduler.alphas_cumprod[timesteps[i - 1]]
214 |                     if i > 0 else self.scheduler.final_alpha_cumprod
215 |                 )
216 | 
217 |                 mu = alpha_prod_t ** 0.5
218 |                 mu_prev = alpha_prod_t_prev ** 0.5
219 |                 sigma = (1 - alpha_prod_t) ** 0.5
220 |                 sigma_prev = (1 - alpha_prod_t_prev) ** 0.5
221 | 
222 |                 eps = self.unet(model_input, t, encoder_hidden_states=cond_batch).sample if self.sd_version != 'ControlNet' \
223 |                     else self.controlnet_pred(x_batch, t, cond_batch, torch.cat([self.canny_cond[b: b + batch_size]]))
224 |                 pred_x0 = (x_batch - sigma_prev * eps) / mu_prev
225 |                 latent_frames[b:b + batch_size] = mu * pred_x0 + sigma * eps
226 | 
227 |             if save_latents and t in timesteps_to_save:
228 |                 torch.save(latent_frames, os.path.join(save_path, 'latents', f'noisy_latents_{t}.pt'))
229 |         torch.save(latent_frames, os.path.join(save_path, 'latents', f'noisy_latents_{t}.pt'))
230 |         return latent_frames
231 | 
232 |     @torch.no_grad()
233 |     def ddim_sample(self, x, cond, batch_size):
234 |         timesteps = self.scheduler.timesteps
235 |         for i, t in enumerate(tqdm(timesteps)):
236 |             for b in range(0, x.shape[0], batch_size):
237 |                 x_batch = x[b:b + batch_size]
238 |                 model_input = x_batch
239 |                 cond_batch = cond.repeat(x_batch.shape[0], 1, 1)
240 |                 
241 |                 if self.sd_version == 'depth':
242 |                     depth_maps = torch.cat([self.depth_maps[b: b + batch_size]])
243 |                     model_input = torch.cat([x_batch, depth_maps],dim=1)
244 |                 
245 |                 alpha_prod_t = self.scheduler.alphas_cumprod[t]
246 |                 alpha_prod_t_prev = (
247 |                     self.scheduler.alphas_cumprod[timesteps[i + 1]]
248 |                     if i < len(timesteps) - 1
249 |                     else self.scheduler.final_alpha_cumprod
250 |                 )
251 |                 mu = alpha_prod_t ** 0.5
252 |                 sigma = (1 - alpha_prod_t) ** 0.5
253 |                 mu_prev = alpha_prod_t_prev ** 0.5
254 |                 sigma_prev = (1 - alpha_prod_t_prev) ** 0.5
255 | 
256 |                 eps = self.unet(model_input, t, encoder_hidden_states=cond_batch).sample if self.sd_version != 'ControlNet' \
257 |                     else self.controlnet_pred(x_batch, t, cond_batch, torch.cat([self.canny_cond[b: b + batch_size]]))
258 | 
259 |                 pred_x0 = (x_batch - sigma * eps) / mu
260 |                 x[b:b + batch_size] = mu_prev * pred_x0 + sigma_prev * eps
261 |         return x
262 | 
263 |     @torch.no_grad()
264 |     def extract_latents(self, 
265 |                         num_steps,
266 |                         save_path,
267 |                         batch_size,
268 |                         timesteps_to_save,
269 |                         inversion_prompt=''):
270 |         self.scheduler.set_timesteps(num_steps)
271 |         cond = self.get_text_embeds(inversion_prompt, "")[1].unsqueeze(0)
272 |         latent_frames = self.latents
273 | 
274 |         inverted_x = self.ddim_inversion(cond,
275 |                                          latent_frames,
276 |                                          save_path,
277 |                                          batch_size=batch_size,
278 |                                          save_latents=True,
279 |                                          timesteps_to_save=timesteps_to_save)
280 |         latent_reconstruction = self.ddim_sample(inverted_x, cond, batch_size=batch_size)
281 |                                                  
282 |         rgb_reconstruction = self.decode_latents(latent_reconstruction)
283 | 
284 |         return rgb_reconstruction
285 | 
286 | 
287 | def prep(opt):
288 |     # timesteps to save
289 |     if opt.sd_version == '2.1':
290 |         model_key = "stabilityai/stable-diffusion-2-1-base"
291 |     elif opt.sd_version == '2.0':
292 |         model_key = "stabilityai/stable-diffusion-2-base"
293 |     elif opt.sd_version == '1.5' or opt.sd_version == 'ControlNet':
294 |         model_key = "runwayml/stable-diffusion-v1-5"
295 |     elif opt.sd_version == 'depth':
296 |         model_key = "stabilityai/stable-diffusion-2-depth"
297 |     toy_scheduler = DDIMScheduler.from_pretrained(model_key, subfolder="scheduler")
298 |     toy_scheduler.set_timesteps(opt.save_steps)
299 |     timesteps_to_save, num_inference_steps = get_timesteps(toy_scheduler, num_inference_steps=opt.save_steps,
300 |                                                            strength=1.0,
301 |                                                            device=device)
302 | 
303 |     seed_everything(1)
304 | 
305 |     save_path = os.path.join(opt.save_dir,
306 |                              f'sd_{opt.sd_version}',
307 |                              Path(opt.data_path).stem,
308 |                              f'steps_{opt.steps}',
309 |                              f'nframes_{opt.n_frames}') 
310 |     os.makedirs(os.path.join(save_path, f'latents'), exist_ok=True)
311 |     add_dict_to_yaml_file(os.path.join(opt.save_dir, 'inversion_prompts.yaml'), Path(opt.data_path).stem, opt.inversion_prompt)    
312 |     # save inversion prompt in a txt file
313 |     with open(os.path.join(save_path, 'inversion_prompt.txt'), 'w') as f:
314 |         f.write(opt.inversion_prompt)
315 |     model = Preprocess(device, opt)
316 |     recon_frames = model.extract_latents(
317 |                                          num_steps=opt.steps,
318 |                                          save_path=save_path,
319 |                                          batch_size=opt.batch_size,
320 |                                          timesteps_to_save=timesteps_to_save,
321 |                                          inversion_prompt=opt.inversion_prompt,
322 |     )
323 | 
324 | 
325 |     if not os.path.isdir(os.path.join(save_path, f'frames')):
326 |         os.mkdir(os.path.join(save_path, f'frames'))
327 |     for i, frame in enumerate(recon_frames):
328 |         T.ToPILImage()(frame).save(os.path.join(save_path, f'frames', f'{i:05d}.png'))
329 |     frames = (recon_frames * 255).to(torch.uint8).cpu().permute(0, 2, 3, 1)
330 |     write_video(os.path.join(save_path, f'inverted.mp4'), frames, fps=10)
331 | 
332 | 
333 | if __name__ == "__main__":
334 |     device = 'cuda'
335 |     parser = argparse.ArgumentParser()
336 |     parser.add_argument('--data_path', type=str,
337 |                         default='data/woman-running.mp4') 
338 |     parser.add_argument('--H', type=int, default=512, 
339 |                         help='for non-square videos, we recommand using 672 x 384 or 384 x 672, aspect ratio 1.75')
340 |     parser.add_argument('--W', type=int, default=512, 
341 |                         help='for non-square videos, we recommand using 672 x 384 or 384 x 672, aspect ratio 1.75')
342 |     parser.add_argument('--save_dir', type=str, default='latents')
343 |     parser.add_argument('--sd_version', type=str, default='2.1', choices=['1.5', '2.0', '2.1', 'ControlNet', 'depth'],
344 |                         help="stable diffusion version")
345 |     parser.add_argument('--steps', type=int, default=500)
346 |     parser.add_argument('--batch_size', type=int, default=40)
347 |     parser.add_argument('--save_steps', type=int, default=50)
348 |     parser.add_argument('--n_frames', type=int, default=40)
349 |     parser.add_argument('--inversion_prompt', type=str, default='a woman running')
350 |     opt = parser.parse_args()
351 |     video_path = opt.data_path
352 |     save_video_frames(video_path, img_size=(opt.W, opt.H))
353 |     opt.data_path = os.path.join('data', Path(video_path).stem)
354 |     prep(opt)
355 | 


--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
 1 | pillow
 2 | diffusers
 3 | ftfy
 4 | transformers
 5 | opencv-python 
 6 | tqdm
 7 | numpy
 8 | pyyaml
 9 | accelerate
10 | xformers
11 | tensorboard
12 | av
13 | kornia
14 | 


--------------------------------------------------------------------------------
/run_tokenflow_pnp.py:
--------------------------------------------------------------------------------
  1 | import glob
  2 | import os
  3 | import numpy as np
  4 | import cv2
  5 | from pathlib import Path
  6 | import torch
  7 | import torch.nn as nn
  8 | import torchvision.transforms as T
  9 | import argparse
 10 | from PIL import Image
 11 | import yaml
 12 | from tqdm import tqdm
 13 | from transformers import logging
 14 | from diffusers import DDIMScheduler, StableDiffusionPipeline
 15 | 
 16 | from tokenflow_utils import *
 17 | from util import save_video, seed_everything
 18 | 
 19 | # suppress partial model loading warning
 20 | logging.set_verbosity_error()
 21 | 
 22 | VAE_BATCH_SIZE = 10
 23 | 
 24 | 
 25 | class TokenFlow(nn.Module):
 26 |     def __init__(self, config):
 27 |         super().__init__()
 28 |         self.config = config
 29 |         self.device = config["device"]
 30 |         
 31 |         sd_version = config["sd_version"]
 32 |         self.sd_version = sd_version
 33 |         if sd_version == '2.1':
 34 |             model_key = "stabilityai/stable-diffusion-2-1-base"
 35 |         elif sd_version == '2.0':
 36 |             model_key = "stabilityai/stable-diffusion-2-base"
 37 |         elif sd_version == '1.5':
 38 |             model_key = "runwayml/stable-diffusion-v1-5"
 39 |         elif sd_version == 'depth':
 40 |             model_key = "stabilityai/stable-diffusion-2-depth"
 41 |         else:
 42 |             raise ValueError(f'Stable-diffusion version {sd_version} not supported.')
 43 | 
 44 |         # Create SD models
 45 |         print('Loading SD model')
 46 | 
 47 |         pipe = StableDiffusionPipeline.from_pretrained(model_key, torch_dtype=torch.float16).to("cuda")
 48 |         # pipe.enable_xformers_memory_efficient_attention()
 49 | 
 50 |         self.vae = pipe.vae
 51 |         self.tokenizer = pipe.tokenizer
 52 |         self.text_encoder = pipe.text_encoder
 53 |         self.unet = pipe.unet
 54 | 
 55 |         self.scheduler = DDIMScheduler.from_pretrained(model_key, subfolder="scheduler")
 56 |         self.scheduler.set_timesteps(config["n_timesteps"], device=self.device)
 57 |         print('SD model loaded')
 58 | 
 59 |         # data
 60 |         self.latents_path = self.get_latents_path()
 61 |         # load frames
 62 |         self.paths, self.frames, self.latents, self.eps = self.get_data()
 63 |         if self.sd_version == 'depth':
 64 |             self.depth_maps = self.prepare_depth_maps()
 65 | 
 66 |         self.text_embeds = self.get_text_embeds(config["prompt"], config["negative_prompt"])
 67 |         pnp_inversion_prompt = self.get_pnp_inversion_prompt()
 68 |         self.pnp_guidance_embeds = self.get_text_embeds(pnp_inversion_prompt, pnp_inversion_prompt).chunk(2)[0]
 69 |     
 70 |     @torch.no_grad()   
 71 |     def prepare_depth_maps(self, model_type='DPT_Large', device='cuda'):
 72 |         depth_maps = []
 73 |         midas = torch.hub.load("intel-isl/MiDaS", model_type)
 74 |         midas.to(device)
 75 |         midas.eval()
 76 | 
 77 |         midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")
 78 | 
 79 |         if model_type == "DPT_Large" or model_type == "DPT_Hybrid":
 80 |             transform = midas_transforms.dpt_transform
 81 |         else:
 82 |             transform = midas_transforms.small_transform
 83 | 
 84 |         for i in range(len(self.paths)):
 85 |             img = cv2.imread(self.paths[i])
 86 |             img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
 87 |             
 88 |             latent_h = img.shape[0] // 8
 89 |             latent_w = img.shape[1] // 8
 90 |             
 91 |             input_batch = transform(img).to(device)
 92 |             prediction = midas(input_batch)
 93 | 
 94 |             depth_map = torch.nn.functional.interpolate(
 95 |                 prediction.unsqueeze(1),
 96 |                 size=(latent_h, latent_w),
 97 |                 mode="bicubic",
 98 |                 align_corners=False,
 99 |             )
100 |             depth_min = torch.amin(depth_map, dim=[1, 2, 3], keepdim=True)
101 |             depth_max = torch.amax(depth_map, dim=[1, 2, 3], keepdim=True)
102 |             depth_map = 2.0 * (depth_map - depth_min) / (depth_max - depth_min) - 1.0
103 |             depth_maps.append(depth_map)
104 | 
105 |         return torch.cat(depth_maps).to(torch.float16).to(self.device)
106 |     
107 |     def get_pnp_inversion_prompt(self):
108 |         inv_prompts_path = os.path.join(str(Path(self.latents_path).parent), 'inversion_prompt.txt')
109 |         # read inversion prompt
110 |         with open(inv_prompts_path, 'r') as f:
111 |             inv_prompt = f.read()
112 |         return inv_prompt
113 | 
114 |     def get_latents_path(self):
115 |         latents_path = os.path.join(config["latents_path"], f'sd_{config["sd_version"]}',
116 |                              Path(config["data_path"]).stem, f'steps_{config["n_inversion_steps"]}')
117 |         latents_path = [x for x in glob.glob(f'{latents_path}/*') if '.' not in Path(x).name]
118 |         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))]
119 |         latents_path = latents_path[np.argmax(n_frames)]
120 |         self.config["n_frames"] = min(max(n_frames), config["n_frames"])
121 |         if self.config["n_frames"] % self.config["batch_size"] != 0:
122 |             # make n_frames divisible by batch_size
123 |             self.config["n_frames"] = self.config["n_frames"] - (self.config["n_frames"] % self.config["batch_size"])
124 |         print("Number of frames: ", self.config["n_frames"])
125 |         return os.path.join(latents_path, 'latents')
126 | 
127 |     @torch.no_grad()
128 |     def get_text_embeds(self, prompt, negative_prompt, batch_size=1):
129 |         # Tokenize text and get embeddings
130 |         text_input = self.tokenizer(prompt, padding='max_length', max_length=self.tokenizer.model_max_length,
131 |                                     truncation=True, return_tensors='pt')
132 |         text_embeddings = self.text_encoder(text_input.input_ids.to(self.device))[0]
133 | 
134 |         # Do the same for unconditional embeddings
135 |         uncond_input = self.tokenizer(negative_prompt, padding='max_length', max_length=self.tokenizer.model_max_length,
136 |                                       return_tensors='pt')
137 | 
138 |         uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(self.device))[0]
139 | 
140 |         # Cat for final embeddings
141 |         text_embeddings = torch.cat([uncond_embeddings] * batch_size + [text_embeddings] * batch_size)
142 |         return text_embeddings
143 | 
144 |     @torch.no_grad()
145 |     def encode_imgs(self, imgs, batch_size=VAE_BATCH_SIZE, deterministic=False):
146 |         imgs = 2 * imgs - 1
147 |         latents = []
148 |         for i in range(0, len(imgs), batch_size):
149 |             posterior = self.vae.encode(imgs[i:i + batch_size]).latent_dist
150 |             latent = posterior.mean if deterministic else posterior.sample()
151 |             latents.append(latent * 0.18215)
152 |         latents = torch.cat(latents)
153 |         return latents
154 | 
155 |     @torch.no_grad()
156 |     def decode_latents(self, latents, batch_size=VAE_BATCH_SIZE):
157 |         latents = 1 / 0.18215 * latents
158 |         imgs = []
159 |         for i in range(0, len(latents), batch_size):
160 |             imgs.append(self.vae.decode(latents[i:i + batch_size]).sample)
161 |         imgs = torch.cat(imgs)
162 |         imgs = (imgs / 2 + 0.5).clamp(0, 1)
163 |         return imgs
164 | 
165 |     
166 |     def get_data(self):
167 |         # load frames
168 |         paths = [os.path.join(config["data_path"], "%05d.jpg" % idx) for idx in
169 |                                range(self.config["n_frames"])]
170 |         if not os.path.exists(paths[0]):
171 |             paths = [os.path.join(config["data_path"], "%05d.png" % idx) for idx in
172 |                                    range(self.config["n_frames"])]
173 |         frames = [Image.open(paths[idx]).convert('RGB') for idx in range(self.config["n_frames"])]
174 |         if frames[0].size[0] == frames[0].size[1]:
175 |             frames = [frame.resize((512, 512), resample=Image.Resampling.LANCZOS) for frame in frames]
176 |         frames = torch.stack([T.ToTensor()(frame) for frame in frames]).to(torch.float16).to(self.device)
177 |         save_video(frames, f'{self.config["output_path"]}/input_fps10.mp4', fps=10)
178 |         save_video(frames, f'{self.config["output_path"]}/input_fps20.mp4', fps=20)
179 |         save_video(frames, f'{self.config["output_path"]}/input_fps30.mp4', fps=30)
180 |         # encode to latents
181 |         latents = self.encode_imgs(frames, deterministic=True).to(torch.float16).to(self.device)
182 |         # get noise
183 |         eps = self.get_ddim_eps(latents, range(self.config["n_frames"])).to(torch.float16).to(self.device)
184 |         return paths, frames, latents, eps
185 | 
186 |     def get_ddim_eps(self, latent, indices):
187 |         noisest = max([int(x.split('_')[-1].split('.')[0]) for x in glob.glob(os.path.join(self.latents_path, f'noisy_latents_*.pt'))])
188 |         latents_path = os.path.join(self.latents_path, f'noisy_latents_{noisest}.pt')
189 |         noisy_latent = torch.load(latents_path)[indices].to(self.device)
190 |         alpha_prod_T = self.scheduler.alphas_cumprod[noisest]
191 |         mu_T, sigma_T = alpha_prod_T ** 0.5, (1 - alpha_prod_T) ** 0.5
192 |         eps = (noisy_latent - mu_T * latent) / sigma_T
193 |         return eps
194 | 
195 |     @torch.no_grad()
196 |     def denoise_step(self, x, t, indices):
197 |         # register the time step and features in pnp injection modules
198 |         source_latents = load_source_latents_t(t, self.latents_path)[indices]
199 |         latent_model_input = torch.cat([source_latents] + ([x] * 2))
200 |         if self.sd_version == 'depth':
201 |             latent_model_input = torch.cat([latent_model_input, torch.cat([self.depth_maps[indices]] * 3)], dim=1)
202 | 
203 |         register_time(self, t.item())
204 | 
205 |         # compute text embeddings
206 |         text_embed_input = torch.cat([self.pnp_guidance_embeds.repeat(len(indices), 1, 1),
207 |                                       torch.repeat_interleave(self.text_embeds, len(indices), dim=0)])
208 | 
209 |         # apply the denoising network
210 |         noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embed_input)['sample']
211 | 
212 |         # perform guidance
213 |         _, noise_pred_uncond, noise_pred_cond = noise_pred.chunk(3)
214 |         noise_pred = noise_pred_uncond + self.config["guidance_scale"] * (noise_pred_cond - noise_pred_uncond)
215 | 
216 |         # compute the denoising step with the reference model
217 |         denoised_latent = self.scheduler.step(noise_pred, t, x)['prev_sample']
218 |         return denoised_latent
219 |     
220 |     @torch.autocast(dtype=torch.float16, device_type='cuda')
221 |     def batched_denoise_step(self, x, t, indices):
222 |         batch_size = self.config["batch_size"]
223 |         denoised_latents = []
224 |         pivotal_idx = torch.randint(batch_size, (len(x)//batch_size,)) + torch.arange(0,len(x),batch_size) 
225 |             
226 |         register_pivotal(self, True)
227 |         self.denoise_step(x[pivotal_idx], t, indices[pivotal_idx])
228 |         register_pivotal(self, False)
229 |         for i, b in enumerate(range(0, len(x), batch_size)):
230 |             register_batch_idx(self, i)
231 |             denoised_latents.append(self.denoise_step(x[b:b + batch_size], t, indices[b:b + batch_size]))
232 |         denoised_latents = torch.cat(denoised_latents)
233 |         return denoised_latents
234 | 
235 |     def init_method(self, conv_injection_t, qk_injection_t):
236 |         self.qk_injection_timesteps = self.scheduler.timesteps[:qk_injection_t] if qk_injection_t >= 0 else []
237 |         self.conv_injection_timesteps = self.scheduler.timesteps[:conv_injection_t] if conv_injection_t >= 0 else []
238 |         register_extended_attention_pnp(self, self.qk_injection_timesteps)
239 |         register_conv_injection(self, self.conv_injection_timesteps)
240 |         set_tokenflow(self.unet)
241 | 
242 |     def save_vae_recon(self):
243 |         os.makedirs(f'{self.config["output_path"]}/vae_recon', exist_ok=True)
244 |         decoded = self.decode_latents(self.latents)
245 |         for i in range(len(decoded)):
246 |             T.ToPILImage()(decoded[i]).save(f'{self.config["output_path"]}/vae_recon/%05d.png' % i)
247 |         save_video(decoded, f'{self.config["output_path"]}/vae_recon_10.mp4', fps=10)
248 |         save_video(decoded, f'{self.config["output_path"]}/vae_recon_20.mp4', fps=20)
249 |         save_video(decoded, f'{self.config["output_path"]}/vae_recon_30.mp4', fps=30)
250 | 
251 |     def edit_video(self):
252 |         os.makedirs(f'{self.config["output_path"]}/img_ode', exist_ok=True)
253 |         self.save_vae_recon()
254 |         pnp_f_t = int(self.config["n_timesteps"] * self.config["pnp_f_t"])
255 |         pnp_attn_t = int(self.config["n_timesteps"] * self.config["pnp_attn_t"])
256 |         self.init_method(conv_injection_t=pnp_f_t, qk_injection_t=pnp_attn_t)
257 |         noisy_latents = self.scheduler.add_noise(self.latents, self.eps, self.scheduler.timesteps[0])
258 |         edited_frames = self.sample_loop(noisy_latents, torch.arange(self.config["n_frames"]))
259 |         save_video(edited_frames, f'{self.config["output_path"]}/tokenflow_PnP_fps_10.mp4')
260 |         save_video(edited_frames, f'{self.config["output_path"]}/tokenflow_PnP_fps_20.mp4', fps=20)
261 |         save_video(edited_frames, f'{self.config["output_path"]}/tokenflow_PnP_fps_30.mp4', fps=30)
262 |         print('Done!')
263 | 
264 |     def sample_loop(self, x, indices):
265 |         os.makedirs(f'{self.config["output_path"]}/img_ode', exist_ok=True)
266 |         for i, t in enumerate(tqdm(self.scheduler.timesteps, desc="Sampling")):
267 |                 x = self.batched_denoise_step(x, t, indices)
268 |         
269 |         decoded_latents = self.decode_latents(x)
270 |         for i in range(len(decoded_latents)):
271 |             T.ToPILImage()(decoded_latents[i]).save(f'{self.config["output_path"]}/img_ode/%05d.png' % i)
272 | 
273 |         return decoded_latents
274 | 
275 | 
276 | def run(config):
277 |     seed_everything(config["seed"])
278 |     print(config)
279 |     editor = TokenFlow(config)
280 |     editor.edit_video()
281 | 
282 | 
283 | if __name__ == '__main__':
284 |     parser = argparse.ArgumentParser()
285 |     parser.add_argument('--config_path', type=str, default='configs/config_pnp.yaml')
286 |     opt = parser.parse_args()
287 |     with open(opt.config_path, "r") as f:
288 |         config = yaml.safe_load(f)
289 |     config["output_path"] = os.path.join(config["output_path"] + f'_pnp_SD_{config["sd_version"]}',
290 |                                              Path(config["data_path"]).stem,
291 |                                              config["prompt"][:240],
292 |                                              f'attn_{config["pnp_attn_t"]}_f_{config["pnp_f_t"]}',
293 |                                              f'batch_size_{str(config["batch_size"])}',
294 |                                              str(config["n_timesteps"]),
295 |     )
296 |     os.makedirs(config["output_path"], exist_ok=True)
297 |     assert os.path.exists(config["data_path"]), "Data path does not exist"
298 |     with open(os.path.join(config["output_path"], "config.yaml"), "w") as f:
299 |         yaml.dump(config, f)
300 |     run(config)
301 | 


--------------------------------------------------------------------------------
/run_tokenflow_sdedit.py:
--------------------------------------------------------------------------------
  1 | import glob
  2 | import os
  3 | from pathlib import Path
  4 | import torch
  5 | import torch.nn as nn
  6 | import torchvision.transforms as T
  7 | import argparse
  8 | from PIL import Image
  9 | import yaml
 10 | from tqdm import tqdm
 11 | from transformers import logging
 12 | from diffusers import DDIMScheduler, StableDiffusionPipeline
 13 | import numpy as np
 14 | 
 15 | from tokenflow_utils import *
 16 | from util import save_video, seed_everything
 17 | 
 18 | # suppress partial model loading warning
 19 | logging.set_verbosity_error()
 20 | 
 21 | VAE_BATCH_SIZE = 10
 22 | 
 23 | 
 24 | # UNET_BATCH_SIZE = 5
 25 | 
 26 | 
 27 | class TokenFlow(nn.Module):
 28 |     def __init__(self, config):
 29 |         super().__init__()
 30 |         self.config = config
 31 |         self.device = config["device"]
 32 |         sd_version = config["sd_version"]
 33 | 
 34 |         if sd_version == '2.1':
 35 |             model_key = "stabilityai/stable-diffusion-2-1-base"
 36 |         elif sd_version == '2.0':
 37 |             model_key = "stabilityai/stable-diffusion-2-base"
 38 |         elif sd_version == '1.5':
 39 |             model_key = "runwayml/stable-diffusion-v1-5"
 40 |         else:
 41 |             raise ValueError(f'Stable-diffusion version {sd_version} not supported.')
 42 | 
 43 |         # Create SD models
 44 |         print('Loading SD model')
 45 | 
 46 |         pipe = StableDiffusionPipeline.from_pretrained(model_key, torch_dtype=torch.float16).to("cuda")
 47 |         # pipe.enable_xformers_memory_efficient_attention()
 48 | 
 49 |         self.vae = pipe.vae
 50 |         self.tokenizer = pipe.tokenizer
 51 |         self.text_encoder = pipe.text_encoder
 52 |         self.unet = pipe.unet
 53 | 
 54 |         self.scheduler = DDIMScheduler.from_pretrained(model_key, subfolder="scheduler")
 55 |         self.scheduler.set_timesteps(config["n_timesteps"], device=self.device)
 56 |         # start from the config["start"] * len(timesteps) timestep
 57 |         self.scheduler.timesteps = self.scheduler.timesteps[int(1 - config["start"] * len(self.scheduler.timesteps)):]
 58 |         print('SD model loaded')
 59 | 
 60 |         # data
 61 |         self.latents_path = self.get_latents_path()
 62 |         self.keyframes_path = [os.path.join(config["data_path"], "%05d.jpg" % idx) for idx in
 63 |                                range(self.config["n_frames"])]
 64 |         if not os.path.exists(self.keyframes_path[0]):
 65 |             self.keyframes_path = [os.path.join(config["data_path"], "%05d.png" % idx) for idx in
 66 |                                    range(self.config["n_frames"])]
 67 |         # load frames
 68 |         self.frames, self.latents, self.eps = self.get_data()
 69 | 
 70 |         self.text_embeds = self.get_text_embeds(config["prompt"], config["negative_prompt"])
 71 |         pnp_inversion_prompt = self.get_pnp_inversion_prompt()
 72 |         self.pnp_guidance_embeds = self.get_text_embeds(pnp_inversion_prompt, pnp_inversion_prompt).chunk(2)[0]
 73 | 
 74 |     def get_pnp_inversion_prompt(self):
 75 |         inv_prompts_path = os.path.join(str(Path(self.latents_path).parent), 'inversion_prompt.txt')
 76 |         # read inversion prompt
 77 |         with open(inv_prompts_path, 'r') as f:
 78 |             inv_prompt = f.read()
 79 |         return inv_prompt
 80 | 
 81 |     def get_latents_path(self):
 82 |         latents_path = os.path.join(config["latents_path"], f'sd_{config["sd_version"]}',
 83 |                              Path(config["data_path"]).stem,)
 84 |         latents_path = [x for x in glob.glob(f'{latents_path}/*/*') if not x.startswith('.')]
 85 |         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))]
 86 |         latents_path = latents_path[np.argmax(n_frames)]
 87 |         self.config["n_frames"] = min(max(n_frames), config["n_frames"])
 88 |         if self.config["n_frames"] % self.config["batch_size"] != 0:
 89 |             # make n_frames divisible by batch_size
 90 |             self.config["n_frames"] = self.config["n_frames"] - (self.config["n_frames"] % self.config["batch_size"])
 91 |         print(self.config["n_frames"])
 92 |         return os.path.join(latents_path, 'latents')
 93 | 
 94 |     @torch.no_grad()
 95 |     def get_text_embeds(self, prompt, negative_prompt, batch_size=1):
 96 |         # Tokenize text and get embeddings
 97 |         text_input = self.tokenizer(prompt, padding='max_length', max_length=self.tokenizer.model_max_length,
 98 |                                     truncation=True, return_tensors='pt')
 99 |         text_embeddings = self.text_encoder(text_input.input_ids.to(self.device))[0]
100 | 
101 |         # Do the same for unconditional embeddings
102 |         uncond_input = self.tokenizer(negative_prompt, padding='max_length', max_length=self.tokenizer.model_max_length,
103 |                                       return_tensors='pt')
104 | 
105 |         uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(self.device))[0]
106 | 
107 |         # Cat for final embeddings
108 |         text_embeddings = torch.cat([uncond_embeddings] * batch_size + [text_embeddings] * batch_size)
109 |         return text_embeddings
110 | 
111 |     @torch.no_grad()
112 |     def encode_imgs(self, imgs, batch_size=VAE_BATCH_SIZE, deterministic=False):
113 |         imgs = 2 * imgs - 1
114 |         latents = []
115 |         for i in range(0, len(imgs), batch_size):
116 |             posterior = self.vae.encode(imgs[i:i + batch_size]).latent_dist
117 |             latent = posterior.mean if deterministic else posterior.sample()
118 |             latents.append(latent * 0.18215)
119 |         latents = torch.cat(latents)
120 |         return latents
121 | 
122 |     @torch.no_grad()
123 |     def decode_latents(self, latents, batch_size=VAE_BATCH_SIZE):
124 |         latents = 1 / 0.18215 * latents
125 |         imgs = []
126 |         for i in range(0, len(latents), batch_size):
127 |             imgs.append(self.vae.decode(latents[i:i + batch_size]).sample)
128 |         imgs = torch.cat(imgs)
129 |         imgs = (imgs / 2 + 0.5).clamp(0, 1)
130 |         return imgs
131 | 
132 |     
133 |     def get_data(self):
134 |         # load frames
135 |         frames = [Image.open(self.keyframes_path[idx]).convert('RGB') for idx in range(self.config["n_frames"])]
136 |         if frames[0].size[0] == frames[0].size[1]:
137 |             frames = [frame.resize((512, 512), resample=Image.Resampling.LANCZOS) for frame in frames]
138 |         frames = torch.stack([T.ToTensor()(frame) for frame in frames]).to(torch.float16).to(self.device)
139 |         # encode to latents
140 |         latents = self.encode_imgs(frames, deterministic=True).to(torch.float16).to(self.device)
141 |         # get noise
142 |         eps = self.get_ddim_eps(latents, range(self.config["n_frames"])).to(torch.float16).to(self.device)
143 |         return frames, latents, eps
144 | 
145 |     def get_ddim_eps(self, latent, indices):
146 |         noisest = max([int(x.split('_')[-1].split('.')[0]) for x in glob.glob(os.path.join(self.latents_path, f'noisy_latents_*.pt'))])
147 |         latents_path = os.path.join(self.latents_path, f'noisy_latents_{noisest}.pt')
148 |         noisy_latent = torch.load(latents_path)[indices].to(self.device)
149 |         alpha_prod_T = self.scheduler.alphas_cumprod[noisest]
150 |         mu_T, sigma_T = alpha_prod_T ** 0.5, (1 - alpha_prod_T) ** 0.5
151 |         eps = (noisy_latent - mu_T * latent) / sigma_T
152 |         return eps
153 | 
154 |     @torch.no_grad()
155 |     def denoise_step(self, x, t, indices):
156 |         # register the time step and features in pnp injection modules
157 |         source_latents = load_source_latents_t(t, self.latents_path)[indices]
158 |         latent_model_input = torch.cat([source_latents] + ([x] * 2))
159 | 
160 |         register_time(self, t.item())
161 | 
162 |         # compute text embeddings
163 |         text_embed_input = torch.cat([self.pnp_guidance_embeds.repeat(len(indices), 1, 1),
164 |                                       torch.repeat_interleave(self.text_embeds, len(indices), dim=0)])
165 | 
166 |         # apply the denoising network
167 |         noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embed_input)['sample']
168 | 
169 |         # perform guidance
170 |         _, noise_pred_uncond, noise_pred_cond = noise_pred.chunk(3)
171 |         noise_pred = noise_pred_uncond + self.config["guidance_scale"] * (noise_pred_cond - noise_pred_uncond)
172 | 
173 |         # compute the denoising step with the reference model
174 |         denoised_latent = self.scheduler.step(noise_pred, t, x)['prev_sample']
175 |         return denoised_latent
176 | 
177 |     @torch.autocast(dtype=torch.float16, device_type='cuda')  
178 |     def batched_denoise_step(self, x, t, indices):
179 |         batch_size = self.config["batch_size"]
180 |         denoised_latents = []
181 |         pivotal_idx = torch.randint(batch_size, (len(x)//batch_size,)) + torch.arange(0,len(x),batch_size) # [int(x) for x in torch.tensor((range(batch_size // 2, len(x) + batch_size // 2, batch_size)))]
182 |         register_pivotal(self, True)
183 |         self.denoise_step(x[pivotal_idx], t, indices[pivotal_idx])
184 |         register_pivotal(self, False)
185 |         for i, b in enumerate(range(0, len(x), batch_size)):
186 |             register_batch_idx(self, i)
187 |             denoised_latents.append(self.denoise_step(x[b:b + batch_size], t, indices[b:b + batch_size]))
188 |         denoised_latents = torch.cat(denoised_latents)
189 |         return denoised_latents
190 | 
191 |     def init_method(self):
192 |         register_extended_attention(self)
193 |         set_tokenflow(self.unet)
194 |         
195 |     def edit_video(self):
196 |         os.makedirs(f'{self.config["output_path"]}/img_ode', exist_ok=True)
197 |         self.init_method()
198 |         noise = self.eps if config["use_ddim_noise"] else torch.randn_like(self.eps[[0]]).repeat(self.eps.shape[0])
199 |         noisy_latents = self.scheduler.add_noise(self.latents, noise, self.scheduler.timesteps[0])
200 |         edited_frames = self.sample_loop(noisy_latents, torch.arange(self.config["n_frames"]))
201 |         save_video(edited_frames, f'{self.config["output_path"]}/tokenflow_SDEdit_fps_10.mp4')
202 |         save_video(edited_frames, f'{self.config["output_path"]}/tokenflow_SDEdit_fps_20.mp4', fps=20)
203 |         save_video(edited_frames, f'{self.config["output_path"]}/tokenflow_SDEdit_fps_30.mp4', fps=30)
204 |         print('Done!')
205 |     
206 | 
207 |     def sample_loop(self, x, indices):
208 |         os.makedirs(f'{self.config["output_path"]}/img_ode', exist_ok=True)
209 |         for i, t in enumerate(tqdm(self.scheduler.timesteps, desc="Sampling")):
210 |                 x = self.batched_denoise_step(x, t, indices)
211 | 
212 |         decoded_latents = self.decode_latents(x)
213 |         for i in range(len(decoded_latents)):
214 |             T.ToPILImage()(decoded_latents[i]).save(f'{self.config["output_path"]}/img_ode/%05d.png' % i)
215 | 
216 |         return decoded_latents
217 |     
218 |     def per_frame_sde(self):
219 |         os.makedirs(f'{self.config["output_path"]}/img_ode', exist_ok=True)
220 |         noisy_latents = self.scheduler.add_noise(self.latents,  self.eps, self.scheduler.timesteps[0])
221 |         edited_frames = self.vanilla_sample_loop(noisy_latents, torch.arange(self.config["n_frames"]))
222 |         save_video(edited_frames, f'{self.config["output_path"]}/vanilla_sde.mp4')
223 |         save_video(edited_frames, f'{self.config["output_path"]}/vanilla_sde_fps20.mp4', fps=20)
224 |         save_video(edited_frames, f'{self.config["output_path"]}/vanilla_sde_fps30.mp4', fps=30)
225 |         print('Done!')
226 | 
227 |     def vanilla_denoise(self, batch, t, text_embed_input):
228 |         latent_model_input = torch.cat(([batch] * 2))
229 |         # apply the denoising network
230 |         noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embed_input)['sample']
231 |         # perform guidance
232 |         noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
233 |         noise_pred = noise_pred_uncond + self.config["guidance_scale"] * (noise_pred_cond - noise_pred_uncond)
234 | 
235 |         # compute the denoising step with the reference model
236 |         batch = self.scheduler.step(noise_pred, t, batch)['prev_sample']
237 |         return batch
238 | 
239 |     def batch_vanilla_denoise_step(self, x, t, text_embed_input):
240 |         denoised = []
241 |         for b in range(0, len(x), self.config["batch_size"]):
242 |             denoised.append(self.vanilla_denoise(x[b:b + self.config["batch_size"]], t, text_embed_input))
243 |         x = torch.cat(denoised)
244 |         return x
245 |     
246 |     @torch.no_grad()
247 |     def vanilla_sample_loop(self, x, indices):
248 |         os.makedirs(f'{self.config["output_path"]}/img_ode_vanilla_sde', exist_ok=True)
249 |         text_embed_input = torch.cat([torch.repeat_interleave(self.text_embeds, config["batch_size"], dim=0)])
250 |         for i, t in enumerate(tqdm(self.scheduler.timesteps, desc="Sampling")):
251 |             x = self.batch_vanilla_denoise_step(x, t, text_embed_input)
252 |             
253 |         decoded_latents = self.decode_latents(x)
254 |         for i in range(len(decoded_latents)):
255 |             T.ToPILImage()(decoded_latents[i]).save(f'{self.config["output_path"]}/img_ode_vanilla_sde/%05d.png' % i)
256 | 
257 |         return decoded_latents
258 | 
259 | 
260 | def run(config):
261 |     seed_everything(config["seed"])
262 |     print(config)
263 |     editor = TokenFlow(config)
264 |     editor.edit_video()
265 | 
266 | 
267 | if __name__ == '__main__':
268 |     parser = argparse.ArgumentParser()
269 |     parser.add_argument('--config_path', type=str, default='configs/config_sdedit.yaml')
270 |     opt = parser.parse_args()
271 |     with open(opt.config_path, "r") as f:
272 |         config = yaml.safe_load(f)
273 | 
274 |     config["output_path"] = os.path.join(config["output_path"] + '_sdedit',
275 |                                              Path(config["data_path"]).stem,
276 |                                              config["prompt"][:240],
277 |                                              f'batch_size_{str(config["batch_size"])}',
278 |                                              str(config["n_timesteps"]) + f'start_{config["start"]}')
279 |     os.makedirs(config["output_path"], exist_ok=True)
280 | 
281 |     with open(os.path.join(config["output_path"], "config.yaml"), "w") as f:
282 |         yaml.dump(config, f)
283 | 
284 |     assert os.path.exists(config["data_path"]), "Data path does not exist"
285 |     run(config)


--------------------------------------------------------------------------------
/tokenflow_utils.py:
--------------------------------------------------------------------------------
  1 | from typing import Type
  2 | import torch
  3 | import os
  4 | 
  5 | from util import isinstance_str, batch_cosine_sim
  6 | 
  7 | def register_pivotal(diffusion_model, is_pivotal):
  8 |     for _, module in diffusion_model.named_modules():
  9 |         # If for some reason this has a different name, create an issue and I'll fix it
 10 |         if isinstance_str(module, "BasicTransformerBlock"):
 11 |             setattr(module, "pivotal_pass", is_pivotal)
 12 |             
 13 | def register_batch_idx(diffusion_model, batch_idx):
 14 |     for _, module in diffusion_model.named_modules():
 15 |         # If for some reason this has a different name, create an issue and I'll fix it
 16 |         if isinstance_str(module, "BasicTransformerBlock"):
 17 |             setattr(module, "batch_idx", batch_idx)
 18 | 
 19 | 
 20 | def register_time(model, t):
 21 |     conv_module = model.unet.up_blocks[1].resnets[1]
 22 |     setattr(conv_module, 't', t)
 23 |     down_res_dict = {0: [0, 1], 1: [0, 1], 2: [0, 1]}
 24 |     up_res_dict = {1: [0, 1, 2], 2: [0, 1, 2], 3: [0, 1, 2]}
 25 |     for res in up_res_dict:
 26 |         for block in up_res_dict[res]:
 27 |             module = model.unet.up_blocks[res].attentions[block].transformer_blocks[0].attn1
 28 |             setattr(module, 't', t)
 29 |             module = model.unet.up_blocks[res].attentions[block].transformer_blocks[0].attn2
 30 |             setattr(module, 't', t)
 31 |     for res in down_res_dict:
 32 |         for block in down_res_dict[res]:
 33 |             module = model.unet.down_blocks[res].attentions[block].transformer_blocks[0].attn1
 34 |             setattr(module, 't', t)
 35 |             module = model.unet.down_blocks[res].attentions[block].transformer_blocks[0].attn2
 36 |             setattr(module, 't', t)
 37 |     module = model.unet.mid_block.attentions[0].transformer_blocks[0].attn1
 38 |     setattr(module, 't', t)
 39 |     module = model.unet.mid_block.attentions[0].transformer_blocks[0].attn2
 40 |     setattr(module, 't', t)
 41 | 
 42 | 
 43 | def load_source_latents_t(t, latents_path):
 44 |     latents_t_path = os.path.join(latents_path, f'noisy_latents_{t}.pt')
 45 |     assert os.path.exists(latents_t_path), f'Missing latents at t {t} path {latents_t_path}'
 46 |     latents = torch.load(latents_t_path)
 47 |     return latents
 48 | 
 49 | def register_conv_injection(model, injection_schedule):
 50 |     def conv_forward(self):
 51 |         def forward(input_tensor, temb):
 52 |             hidden_states = input_tensor
 53 | 
 54 |             hidden_states = self.norm1(hidden_states)
 55 |             hidden_states = self.nonlinearity(hidden_states)
 56 | 
 57 |             if self.upsample is not None:
 58 |                 # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
 59 |                 if hidden_states.shape[0] >= 64:
 60 |                     input_tensor = input_tensor.contiguous()
 61 |                     hidden_states = hidden_states.contiguous()
 62 |                 input_tensor = self.upsample(input_tensor)
 63 |                 hidden_states = self.upsample(hidden_states)
 64 |             elif self.downsample is not None:
 65 |                 input_tensor = self.downsample(input_tensor)
 66 |                 hidden_states = self.downsample(hidden_states)
 67 | 
 68 |             hidden_states = self.conv1(hidden_states)
 69 | 
 70 |             if temb is not None:
 71 |                 temb = self.time_emb_proj(self.nonlinearity(temb))[:, :, None, None]
 72 | 
 73 |             if temb is not None and self.time_embedding_norm == "default":
 74 |                 hidden_states = hidden_states + temb
 75 | 
 76 |             hidden_states = self.norm2(hidden_states)
 77 | 
 78 |             if temb is not None and self.time_embedding_norm == "scale_shift":
 79 |                 scale, shift = torch.chunk(temb, 2, dim=1)
 80 |                 hidden_states = hidden_states * (1 + scale) + shift
 81 | 
 82 |             hidden_states = self.nonlinearity(hidden_states)
 83 | 
 84 |             hidden_states = self.dropout(hidden_states)
 85 |             hidden_states = self.conv2(hidden_states)
 86 |             if self.injection_schedule is not None and (self.t in self.injection_schedule or self.t == 1000):
 87 |                 source_batch_size = int(hidden_states.shape[0] // 3)
 88 |                 # inject unconditional
 89 |                 hidden_states[source_batch_size:2 * source_batch_size] = hidden_states[:source_batch_size]
 90 |                 # inject conditional
 91 |                 hidden_states[2 * source_batch_size:] = hidden_states[:source_batch_size]
 92 | 
 93 |             if self.conv_shortcut is not None:
 94 |                 input_tensor = self.conv_shortcut(input_tensor)
 95 | 
 96 |             output_tensor = (input_tensor + hidden_states) / self.output_scale_factor
 97 | 
 98 |             return output_tensor
 99 | 
100 |         return forward
101 | 
102 |     conv_module = model.unet.up_blocks[1].resnets[1]
103 |     conv_module.forward = conv_forward(conv_module)
104 |     setattr(conv_module, 'injection_schedule', injection_schedule)
105 | 
106 | def register_extended_attention_pnp(model, injection_schedule):
107 |     def sa_forward(self):
108 |         to_out = self.to_out
109 |         if type(to_out) is torch.nn.modules.container.ModuleList:
110 |             to_out = self.to_out[0]
111 |         else:
112 |             to_out = self.to_out
113 | 
114 |         def forward(x, encoder_hidden_states=None, attention_mask=None):
115 |             batch_size, sequence_length, dim = x.shape
116 |             h = self.heads
117 |             n_frames = batch_size // 3
118 |             is_cross = encoder_hidden_states is not None
119 |             encoder_hidden_states = encoder_hidden_states if is_cross else x
120 |             q = self.to_q(x)
121 |             k = self.to_k(encoder_hidden_states)
122 |             v = self.to_v(encoder_hidden_states)
123 | 
124 |             if self.injection_schedule is not None and (self.t in self.injection_schedule or self.t == 1000):
125 |                 # inject unconditional
126 |                 q[n_frames:2 * n_frames] = q[:n_frames]
127 |                 k[n_frames:2 * n_frames] = k[:n_frames]
128 |                 # inject conditional
129 |                 q[2 * n_frames:] = q[:n_frames]
130 |                 k[2 * n_frames:] = k[:n_frames]
131 | 
132 |             k_source = k[:n_frames]
133 |             k_uncond = k[n_frames:2 * n_frames].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1)
134 |             k_cond = k[2 * n_frames:].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1)
135 | 
136 |             v_source = v[:n_frames]
137 |             v_uncond = v[n_frames:2 * n_frames].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1)
138 |             v_cond = v[2 * n_frames:].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1)
139 | 
140 |             q_source = self.head_to_batch_dim(q[:n_frames])
141 |             q_uncond = self.head_to_batch_dim(q[n_frames:2 * n_frames])
142 |             q_cond = self.head_to_batch_dim(q[2 * n_frames:])
143 |             k_source = self.head_to_batch_dim(k_source)
144 |             k_uncond = self.head_to_batch_dim(k_uncond)
145 |             k_cond = self.head_to_batch_dim(k_cond)
146 |             v_source = self.head_to_batch_dim(v_source)
147 |             v_uncond = self.head_to_batch_dim(v_uncond)
148 |             v_cond = self.head_to_batch_dim(v_cond)
149 | 
150 | 
151 |             q_src = q_source.view(n_frames, h, sequence_length, dim // h)
152 |             k_src = k_source.view(n_frames, h, sequence_length, dim // h)
153 |             v_src = v_source.view(n_frames, h, sequence_length, dim // h)
154 |             q_uncond = q_uncond.view(n_frames, h, sequence_length, dim // h)
155 |             k_uncond = k_uncond.view(n_frames, h, sequence_length * n_frames, dim // h)
156 |             v_uncond = v_uncond.view(n_frames, h, sequence_length * n_frames, dim // h)
157 |             q_cond = q_cond.view(n_frames, h, sequence_length, dim // h)
158 |             k_cond = k_cond.view(n_frames, h, sequence_length * n_frames, dim // h)
159 |             v_cond = v_cond.view(n_frames, h, sequence_length * n_frames, dim // h)
160 | 
161 |             out_source_all = []
162 |             out_uncond_all = []
163 |             out_cond_all = []
164 |             
165 |             single_batch = n_frames<=12
166 |             b = n_frames if single_batch else 1
167 | 
168 |             for frame in range(0, n_frames, b):
169 |                 out_source = []
170 |                 out_uncond = []
171 |                 out_cond = []
172 |                 for j in range(h):
173 |                     sim_source_b = torch.bmm(q_src[frame: frame+ b, j], k_src[frame: frame+ b, j].transpose(-1, -2)) * self.scale
174 |                     sim_uncond_b = torch.bmm(q_uncond[frame: frame+ b, j], k_uncond[frame: frame+ b, j].transpose(-1, -2)) * self.scale
175 |                     sim_cond = torch.bmm(q_cond[frame: frame+ b, j], k_cond[frame: frame+ b, j].transpose(-1, -2)) * self.scale
176 | 
177 |                     out_source.append(torch.bmm(sim_source_b.softmax(dim=-1), v_src[frame: frame+ b, j]))
178 |                     out_uncond.append(torch.bmm(sim_uncond_b.softmax(dim=-1), v_uncond[frame: frame+ b, j]))
179 |                     out_cond.append(torch.bmm(sim_cond.softmax(dim=-1), v_cond[frame: frame+ b, j]))
180 | 
181 |                 out_source = torch.cat(out_source, dim=0)
182 |                 out_uncond = torch.cat(out_uncond, dim=0) 
183 |                 out_cond = torch.cat(out_cond, dim=0) 
184 |                 if single_batch:
185 |                     out_source = out_source.view(h, n_frames,sequence_length, dim // h).permute(1, 0, 2, 3).reshape(h * n_frames, sequence_length, -1)
186 |                     out_uncond = out_uncond.view(h, n_frames,sequence_length, dim // h).permute(1, 0, 2, 3).reshape(h * n_frames, sequence_length, -1)
187 |                     out_cond = out_cond.view(h, n_frames,sequence_length, dim // h).permute(1, 0, 2, 3).reshape(h * n_frames, sequence_length, -1)
188 |                 out_source_all.append(out_source)
189 |                 out_uncond_all.append(out_uncond)
190 |                 out_cond_all.append(out_cond)
191 |             
192 |             out_source = torch.cat(out_source_all, dim=0)
193 |             out_uncond = torch.cat(out_uncond_all, dim=0)
194 |             out_cond = torch.cat(out_cond_all, dim=0)
195 |                 
196 |             out = torch.cat([out_source, out_uncond, out_cond], dim=0)
197 |             out = self.batch_to_head_dim(out)
198 | 
199 |             return to_out(out)
200 | 
201 |         return forward
202 | 
203 |     for _, module in model.unet.named_modules():
204 |         if isinstance_str(module, "BasicTransformerBlock"):
205 |             module.attn1.forward = sa_forward(module.attn1)
206 |             setattr(module.attn1, 'injection_schedule', [])
207 | 
208 |     res_dict = {1: [1, 2], 2: [0, 1, 2], 3: [0, 1, 2]}
209 |     # we are injecting attention in blocks 4 - 11 of the decoder, so not in the first block of the lowest resolution
210 |     for res in res_dict:
211 |         for block in res_dict[res]:
212 |             module = model.unet.up_blocks[res].attentions[block].transformer_blocks[0].attn1
213 |             module.forward = sa_forward(module)
214 |             setattr(module, 'injection_schedule', injection_schedule)
215 | 
216 | def register_extended_attention(model):
217 |     def sa_forward(self):
218 |         to_out = self.to_out
219 |         if type(to_out) is torch.nn.modules.container.ModuleList:
220 |             to_out = self.to_out[0]
221 |         else:
222 |             to_out = self.to_out
223 | 
224 |         def forward(x, encoder_hidden_states=None, attention_mask=None):
225 |             batch_size, sequence_length, dim = x.shape
226 |             h = self.heads
227 |             n_frames = batch_size // 3
228 |             is_cross = encoder_hidden_states is not None
229 |             encoder_hidden_states = encoder_hidden_states if is_cross else x
230 |             q = self.to_q(x)
231 |             k = self.to_k(encoder_hidden_states)
232 |             v = self.to_v(encoder_hidden_states)
233 | 
234 |             k_source = k[:n_frames]
235 |             k_uncond = k[n_frames: 2*n_frames].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1)
236 |             k_cond = k[2*n_frames:].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1)
237 |             v_source = v[:n_frames]
238 |             v_uncond = v[n_frames:2*n_frames].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1)
239 |             v_cond = v[2*n_frames:].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1)
240 | 
241 |             q_source = self.head_to_batch_dim(q[:n_frames])
242 |             q_uncond = self.head_to_batch_dim(q[n_frames: 2*n_frames])
243 |             q_cond = self.head_to_batch_dim(q[2 * n_frames:])
244 |             k_source = self.head_to_batch_dim(k_source)
245 |             k_uncond = self.head_to_batch_dim(k_uncond)
246 |             k_cond = self.head_to_batch_dim(k_cond)
247 |             v_source = self.head_to_batch_dim(v_source)
248 |             v_uncond = self.head_to_batch_dim(v_uncond)
249 |             v_cond = self.head_to_batch_dim(v_cond)
250 | 
251 |             out_source = []
252 |             out_uncond = []
253 |             out_cond = []
254 | 
255 |             q_src = q_source.view(n_frames, h, sequence_length, dim // h)
256 |             k_src = k_source.view(n_frames, h, sequence_length, dim // h)
257 |             v_src = v_source.view(n_frames, h, sequence_length, dim // h)
258 |             q_uncond = q_uncond.view(n_frames, h, sequence_length, dim // h)
259 |             k_uncond = k_uncond.view(n_frames, h, sequence_length * n_frames, dim // h)
260 |             v_uncond = v_uncond.view(n_frames, h, sequence_length * n_frames, dim // h)
261 |             q_cond = q_cond.view(n_frames, h, sequence_length, dim // h)
262 |             k_cond = k_cond.view(n_frames, h, sequence_length * n_frames, dim // h)
263 |             v_cond = v_cond.view(n_frames, h, sequence_length * n_frames, dim // h)
264 | 
265 |             for j in range(h):
266 |                 sim_source_b = torch.bmm(q_src[:, j], k_src[:, j].transpose(-1, -2)) * self.scale
267 |                 sim_uncond_b = torch.bmm(q_uncond[:, j], k_uncond[:, j].transpose(-1, -2)) * self.scale
268 |                 sim_cond = torch.bmm(q_cond[:, j], k_cond[:, j].transpose(-1, -2)) * self.scale
269 | 
270 |                 out_source.append(torch.bmm(sim_source_b.softmax(dim=-1), v_src[:, j]))
271 |                 out_uncond.append(torch.bmm(sim_uncond_b.softmax(dim=-1), v_uncond[:, j]))
272 |                 out_cond.append(torch.bmm(sim_cond.softmax(dim=-1), v_cond[:, j]))
273 | 
274 |             out_source = torch.cat(out_source, dim=0).view(h, n_frames,sequence_length, dim // h).permute(1, 0, 2, 3).reshape(h * n_frames, sequence_length, -1)
275 |             out_uncond = torch.cat(out_uncond, dim=0).view(h, n_frames,sequence_length, dim // h).permute(1, 0, 2, 3).reshape(h * n_frames, sequence_length, -1)
276 |             out_cond = torch.cat(out_cond, dim=0).view(h, n_frames,sequence_length, dim // h).permute(1, 0, 2, 3).reshape(h * n_frames, sequence_length, -1)
277 | 
278 |             out = torch.cat([out_source, out_uncond, out_cond], dim=0)
279 |             out = self.batch_to_head_dim(out)
280 | 
281 |             return to_out(out)
282 | 
283 |         return forward
284 | 
285 |     for _, module in model.unet.named_modules():
286 |         if isinstance_str(module, "BasicTransformerBlock"):
287 |             module.attn1.forward = sa_forward(module.attn1)
288 | 
289 |     res_dict = {1: [1, 2], 2: [0, 1, 2], 3: [0, 1, 2]}
290 |     # we are injecting attention in blocks 4 - 11 of the decoder, so not in the first block of the lowest resolution
291 |     for res in res_dict:
292 |         for block in res_dict[res]:
293 |             module = model.unet.up_blocks[res].attentions[block].transformer_blocks[0].attn1
294 |             module.forward = sa_forward(module)
295 | 
296 | def make_tokenflow_attention_block(block_class: Type[torch.nn.Module]) -> Type[torch.nn.Module]:
297 | 
298 |     class TokenFlowBlock(block_class):
299 | 
300 |         def forward(
301 |             self,
302 |             hidden_states,
303 |             attention_mask=None,
304 |             encoder_hidden_states=None,
305 |             encoder_attention_mask=None,
306 |             timestep=None,
307 |             cross_attention_kwargs=None,
308 |             class_labels=None,
309 |         ) -> torch.Tensor:
310 |             
311 |             batch_size, sequence_length, dim = hidden_states.shape
312 |             n_frames = batch_size // 3
313 |             mid_idx = n_frames // 2
314 |             hidden_states = hidden_states.view(3, n_frames, sequence_length, dim)
315 | 
316 |             if self.use_ada_layer_norm:
317 |                 norm_hidden_states = self.norm1(hidden_states, timestep)
318 |             elif self.use_ada_layer_norm_zero:
319 |                 norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
320 |                     hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype
321 |                 )
322 |             else:
323 |                 norm_hidden_states = self.norm1(hidden_states)
324 | 
325 |             norm_hidden_states = norm_hidden_states.view(3, n_frames, sequence_length, dim)
326 |             if self.pivotal_pass:
327 |                 self.pivot_hidden_states = norm_hidden_states
328 |             else:
329 |                 idx1 = []
330 |                 idx2 = [] 
331 |                 batch_idxs = [self.batch_idx]
332 |                 if self.batch_idx > 0:
333 |                     batch_idxs.append(self.batch_idx - 1)
334 |                 
335 |                 sim = batch_cosine_sim(norm_hidden_states[0].reshape(-1, dim),
336 |                                         self.pivot_hidden_states[0][batch_idxs].reshape(-1, dim))
337 |                 if len(batch_idxs) == 2:
338 |                     sim1, sim2 = sim.chunk(2, dim=1)
339 |                     # sim: n_frames * seq_len, len(batch_idxs) * seq_len
340 |                     idx1.append(sim1.argmax(dim=-1))  # n_frames * seq_len
341 |                     idx2.append(sim2.argmax(dim=-1))  # n_frames * seq_len
342 |                 else:
343 |                     idx1.append(sim.argmax(dim=-1))
344 |                 idx1 = torch.stack(idx1 * 3, dim=0) # 3, n_frames * seq_len
345 |                 idx1 = idx1.squeeze(1)
346 |                 if len(batch_idxs) == 2:
347 |                     idx2 = torch.stack(idx2 * 3, dim=0) # 3, n_frames * seq_len
348 |                     idx2 = idx2.squeeze(1)
349 | 
350 |             # 1. Self-Attention
351 |             cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
352 |             if self.pivotal_pass:
353 |                 # norm_hidden_states.shape = 3, n_frames * seq_len, dim
354 |                 self.attn_output = self.attn1(
355 |                         norm_hidden_states.view(batch_size, sequence_length, dim),
356 |                         encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,
357 |                         **cross_attention_kwargs,
358 |                     )
359 |                 # 3, n_frames * seq_len, dim - > 3 * n_frames, seq_len, dim
360 |                 self.kf_attn_output = self.attn_output 
361 |             else:
362 |                 batch_kf_size, _, _ = self.kf_attn_output.shape
363 |                 self.attn_output = self.kf_attn_output.view(3, batch_kf_size // 3, sequence_length, dim)[:,
364 |                                    batch_idxs]  # 3, n_frames, seq_len, dim --> 3, len(batch_idxs), seq_len, dim
365 |             if self.use_ada_layer_norm_zero:
366 |                 self.attn_output = gate_msa.unsqueeze(1) * self.attn_output
367 | 
368 |             # gather values from attn_output, using idx as indices, and get a tensor of shape 3, n_frames, seq_len, dim
369 |             if not self.pivotal_pass:
370 |                 if len(batch_idxs) == 2:
371 |                     attn_1, attn_2 = self.attn_output[:, 0], self.attn_output[:, 1]
372 |                     attn_output1 = attn_1.gather(dim=1, index=idx1.unsqueeze(-1).repeat(1, 1, dim))
373 |                     attn_output2 = attn_2.gather(dim=1, index=idx2.unsqueeze(-1).repeat(1, 1, dim))
374 | 
375 |                     s = torch.arange(0, n_frames).to(idx1.device) + batch_idxs[0] * n_frames
376 |                     # distance from the pivot
377 |                     p1 = batch_idxs[0] * n_frames + n_frames // 2
378 |                     p2 = batch_idxs[1] * n_frames + n_frames // 2
379 |                     d1 = torch.abs(s - p1)
380 |                     d2 = torch.abs(s - p2)
381 |                     # weight
382 |                     w1 = d2 / (d1 + d2)
383 |                     w1 = torch.sigmoid(w1)
384 |                     
385 |                     w1 = w1.unsqueeze(0).unsqueeze(-1).unsqueeze(-1).repeat(3, 1, sequence_length, dim)
386 |                     attn_output1 = attn_output1.view(3, n_frames, sequence_length, dim)
387 |                     attn_output2 = attn_output2.view(3, n_frames, sequence_length, dim)
388 |                     attn_output = w1 * attn_output1 + (1 - w1) * attn_output2
389 |                 else:
390 |                     attn_output = self.attn_output[:,0].gather(dim=1, index=idx1.unsqueeze(-1).repeat(1, 1, dim))
391 | 
392 |                 attn_output = attn_output.reshape(
393 |                         batch_size, sequence_length, dim)  # 3 * n_frames, seq_len, dim
394 |             else:
395 |                 attn_output = self.attn_output
396 |             hidden_states = hidden_states.reshape(batch_size, sequence_length, dim)  # 3 * n_frames, seq_len, dim
397 |             hidden_states = attn_output + hidden_states
398 | 
399 |             if self.attn2 is not None:
400 |                 norm_hidden_states = (
401 |                     self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)
402 |                 )
403 | 
404 |                 # 2. Cross-Attention
405 |                 attn_output = self.attn2(
406 |                     norm_hidden_states,
407 |                     encoder_hidden_states=encoder_hidden_states,
408 |                     attention_mask=encoder_attention_mask,
409 |                     **cross_attention_kwargs,
410 |                 )
411 |                 hidden_states = attn_output + hidden_states
412 | 
413 |             # 3. Feed-forward
414 |             norm_hidden_states = self.norm3(hidden_states)
415 | 
416 |             if self.use_ada_layer_norm_zero:
417 |                 norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
418 | 
419 | 
420 |             ff_output = self.ff(norm_hidden_states)
421 | 
422 |             if self.use_ada_layer_norm_zero:
423 |                 ff_output = gate_mlp.unsqueeze(1) * ff_output
424 | 
425 |             hidden_states = ff_output + hidden_states
426 | 
427 |             return hidden_states
428 | 
429 |     return TokenFlowBlock
430 | 
431 | 
432 | def set_tokenflow(
433 |         model: torch.nn.Module):
434 |     """
435 |     Sets the tokenflow attention blocks in a model.
436 |     """
437 | 
438 |     for _, module in model.named_modules():
439 |         if isinstance_str(module, "BasicTransformerBlock"):
440 |             make_tokenflow_block_fn = make_tokenflow_attention_block 
441 |             module.__class__ = make_tokenflow_block_fn(module.__class__)
442 | 
443 |             # Something needed for older versions of diffusers
444 |             if not hasattr(module, "use_ada_layer_norm_zero"):
445 |                 module.use_ada_layer_norm = False
446 |                 module.use_ada_layer_norm_zero = False
447 | 
448 |     return model
449 | 


--------------------------------------------------------------------------------
/util.py:
--------------------------------------------------------------------------------
  1 | from pathlib import Path
  2 | from PIL import Image
  3 | import torch
  4 | import yaml
  5 | import math
  6 | 
  7 | import torchvision.transforms as T
  8 | from torchvision.io import read_video,write_video
  9 | import os
 10 | import random
 11 | import numpy as np
 12 | from torchvision.io import write_video
 13 | # from kornia.filters import joint_bilateral_blur
 14 | from kornia.geometry.transform import remap
 15 | from kornia.utils.grid import create_meshgrid
 16 | import cv2
 17 | 
 18 | def save_video_frames(video_path, img_size=(512,512)):
 19 |     video, _, _ = read_video(video_path, output_format="TCHW")
 20 |     # rotate video -90 degree if video is .mov format. this is a weird bug in torchvision
 21 |     if video_path.endswith('.mov'):
 22 |         video = T.functional.rotate(video, -90)
 23 |     video_name = Path(video_path).stem
 24 |     os.makedirs(f'data/{video_name}', exist_ok=True)
 25 |     for i in range(len(video)):
 26 |         ind = str(i).zfill(5)
 27 |         image = T.ToPILImage()(video[i])
 28 |         image_resized = image.resize((img_size),  resample=Image.Resampling.LANCZOS)
 29 |         image_resized.save(f'data/{video_name}/{ind}.png')
 30 | 
 31 | def add_dict_to_yaml_file(file_path, key, value):
 32 |     data = {}
 33 | 
 34 |     # If the file already exists, load its contents into the data dictionary
 35 |     if os.path.exists(file_path):
 36 |         with open(file_path, 'r') as file:
 37 |             data = yaml.safe_load(file)
 38 | 
 39 |     # Add or update the key-value pair
 40 |     data[key] = value
 41 | 
 42 |     # Save the data back to the YAML file
 43 |     with open(file_path, 'w') as file:
 44 |         yaml.dump(data, file)
 45 |         
 46 | def isinstance_str(x: object, cls_name: str):
 47 |     """
 48 |     Checks whether x has any class *named* cls_name in its ancestry.
 49 |     Doesn't require access to the class's implementation.
 50 |     
 51 |     Useful for patching!
 52 |     """
 53 | 
 54 |     for _cls in x.__class__.__mro__:
 55 |         if _cls.__name__ == cls_name:
 56 |             return True
 57 |     
 58 |     return False
 59 | 
 60 | 
 61 | def batch_cosine_sim(x, y):
 62 |     if type(x) is list:
 63 |         x = torch.cat(x, dim=0)
 64 |     if type(y) is list:
 65 |         y = torch.cat(y, dim=0)
 66 |     x = x / x.norm(dim=-1, keepdim=True)
 67 |     y = y / y.norm(dim=-1, keepdim=True)
 68 |     similarity = x @ y.T
 69 |     return similarity
 70 | 
 71 | 
 72 | def load_imgs(data_path, n_frames, device='cuda', pil=False):
 73 |     imgs = []
 74 |     pils = []
 75 |     for i in range(n_frames):
 76 |         img_path = os.path.join(data_path, "%05d.jpg" % i)
 77 |         if not os.path.exists(img_path):
 78 |             img_path = os.path.join(data_path, "%05d.png" % i)
 79 |         img_pil = Image.open(img_path)
 80 |         pils.append(img_pil)
 81 |         img = T.ToTensor()(img_pil).unsqueeze(0)
 82 |         imgs.append(img)
 83 |     if pil:
 84 |         return torch.cat(imgs).to(device), pils
 85 |     return torch.cat(imgs).to(device)
 86 | 
 87 | 
 88 | def save_video(raw_frames, save_path, fps=10):
 89 |     video_codec = "libx264"
 90 |     video_options = {
 91 |         "crf": "18",  # Constant Rate Factor (lower value = higher quality, 18 is a good balance)
 92 |         "preset": "slow",  # Encoding preset (e.g., ultrafast, superfast, veryfast, faster, fast, medium, slow, slower, veryslow)
 93 |     }
 94 | 
 95 |     frames = (raw_frames * 255).to(torch.uint8).cpu().permute(0, 2, 3, 1)
 96 |     write_video(save_path, frames, fps=fps, video_codec=video_codec, options=video_options)
 97 | 
 98 | 
 99 | def seed_everything(seed):
100 |     torch.manual_seed(seed)
101 |     torch.cuda.manual_seed(seed)
102 |     random.seed(seed)
103 |     np.random.seed(seed)
104 | 
105 | 
106 | 


--------------------------------------------------------------------------------