Duplicate from weizmannscience/tokenflow

Co-authored-by: Linoy Tsaban <LinoyTsaban@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,37 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/woman-running.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/running_dog.mp4 filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,13 @@
+---
+title: Tokenflow
+emoji: 🐠
+colorFrom: purple
+colorTo: blue
+sdk: gradio
+sdk_version: 3.41.2
+app_file: app.py
+pinned: false
+duplicated_from: weizmannscience/tokenflow
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,374 @@
+import gradio as gr
+import torch
+from diffusers import StableDiffusionPipeline, DDIMScheduler
+from utils import video_to_frames, add_dict_to_yaml_file, save_video, seed_everything
+# from diffusers.utils import export_to_video
+from tokenflow_pnp import TokenFlow
+from preprocess_utils import *
+from tokenflow_utils import *
+# load sd model
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model_id = "stabilityai/stable-diffusion-2-1-base"
+
+# components for the Preprocessor
+scheduler = DDIMScheduler.from_pretrained(model_id, subfolder="scheduler")
+vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae", revision="fp16",
+                                                 torch_dtype=torch.float16).to(device)
+tokenizer = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer")
+text_encoder = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder", revision="fp16",
+                                                  torch_dtype=torch.float16).to(device)
+unet = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet", revision="fp16",
+                                           torch_dtype=torch.float16).to(device)
+
+# pipe for TokenFlow
+tokenflow_pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")
+tokenflow_pipe.enable_xformers_memory_efficient_attention()
+
+def randomize_seed_fn():
+    seed = random.randint(0, np.iinfo(np.int32).max)
+    return seed
+    
+def reset_do_inversion():
+    return True
+
+def get_example():
+    case = [
+        [
+            'examples/wolf.mp4',     
+        ],
+        [
+            'examples/woman-running.mp4',     
+        ],
+        [
+            'examples/cutting_bread.mp4',
+        ],
+        [
+            'examples/running_dog.mp4',
+        ]
+    ]
+    return case
+
+
+def prep(config):
+    # timesteps to save
+    if config["sd_version"] == '2.1':
+        model_key = "stabilityai/stable-diffusion-2-1-base"
+    elif config["sd_version"] == '2.0':
+        model_key = "stabilityai/stable-diffusion-2-base"
+    elif config["sd_version"] == '1.5' or config["sd_version"] == 'ControlNet':
+        model_key = "runwayml/stable-diffusion-v1-5"
+    elif config["sd_version"] == 'depth':
+        model_key = "stabilityai/stable-diffusion-2-depth"
+    toy_scheduler = DDIMScheduler.from_pretrained(model_key, subfolder="scheduler")
+    toy_scheduler.set_timesteps(config["save_steps"])
+    print("config[save_steps]", config["save_steps"])
+    timesteps_to_save, num_inference_steps = get_timesteps(toy_scheduler, num_inference_steps=config["save_steps"],
+                                                           strength=1.0,
+                                                           device=device)
+    print("YOOOO timesteps to save", timesteps_to_save)
+
+    # seed_everything(config["seed"])
+    if not config["frames"]: # original non demo setting
+        save_path = os.path.join(config["save_dir"],
+                                 f'sd_{config["sd_version"]}',
+                                 Path(config["data_path"]).stem,
+                                 f'steps_{config["steps"]}',
+                                 f'nframes_{config["n_frames"]}') 
+        os.makedirs(os.path.join(save_path, f'latents'), exist_ok=True)
+        add_dict_to_yaml_file(os.path.join(config["save_dir"], 'inversion_prompts.yaml'), Path(config["data_path"]).stem, config["inversion_prompt"])    
+        # save inversion prompt in a txt file
+        with open(os.path.join(save_path, 'inversion_prompt.txt'), 'w') as f:
+            f.write(config["inversion_prompt"])
+    else:
+        save_path = None
+    
+    model = Preprocess(device, config,
+                      vae=vae,
+                      text_encoder=text_encoder,
+                      scheduler=scheduler,
+                      tokenizer=tokenizer,
+                      unet=unet)
+    print(type(model.config["batch_size"]))
+    frames, latents, total_inverted_latents, rgb_reconstruction = model.extract_latents(
+                                         num_steps=model.config["steps"],
+                                         save_path=save_path,
+                                         batch_size=model.config["batch_size"],
+                                         timesteps_to_save=timesteps_to_save,
+                                         inversion_prompt=model.config["inversion_prompt"],
+    )
+
+    
+    return frames, latents, total_inverted_latents, rgb_reconstruction
+    
+def preprocess_and_invert(input_video,
+                          frames,
+                          latents,
+                          inverted_latents,
+                          seed, 
+                          randomize_seed,
+                          do_inversion,
+                          # save_dir: str = "latents",
+                          steps,
+                          n_timesteps = 50,
+                          batch_size: int = 8,
+                          n_frames: int = 40,
+                          inversion_prompt:str = '',
+                          
+              ):
+    sd_version = "2.1"
+    height = 512
+    weidth: int = 512
+    print("n timesteps", n_timesteps)
+    if do_inversion or randomize_seed:
+        preprocess_config = {}
+        preprocess_config['H'] = height
+        preprocess_config['W'] = weidth
+        preprocess_config['save_dir'] = 'latents'
+        preprocess_config['sd_version'] = sd_version
+        preprocess_config['steps'] = steps
+        preprocess_config['batch_size'] = batch_size
+        preprocess_config['save_steps'] = int(n_timesteps)
+        preprocess_config['n_frames'] = n_frames
+        preprocess_config['seed'] = seed
+        preprocess_config['inversion_prompt'] = inversion_prompt
+        preprocess_config['frames'] = video_to_frames(input_video)
+        preprocess_config['data_path'] = input_video.split(".")[0]
+        
+
+        if randomize_seed:
+            seed = randomize_seed_fn()
+        seed_everything(seed)
+        
+        frames, latents, total_inverted_latents, rgb_reconstruction = prep(preprocess_config)
+        print(total_inverted_latents.keys())
+        print(len(total_inverted_latents.keys()))
+        frames = gr.State(value=frames)
+        latents = gr.State(value=latents)
+        inverted_latents = gr.State(value=total_inverted_latents)
+        do_inversion = False
+   
+    return frames, latents, inverted_latents, do_inversion
+
+
+def edit_with_pnp(input_video,
+                  frames, 
+                  latents,
+                  inverted_latents,
+                  seed,
+                  randomize_seed,
+                  do_inversion,
+                  steps,
+                  prompt: str = "a marble sculpture of a woman running, Venus de Milo",
+                  # negative_prompt: str = "ugly, blurry, low res, unrealistic, unaesthetic",
+                  pnp_attn_t: float = 0.5,
+                  pnp_f_t: float = 0.8,
+                  batch_size: int = 8, #needs to be the same as for preprocess
+                  n_frames: int = 40,#needs to be the same as for preprocess
+                  n_timesteps: int = 50,
+                  gudiance_scale: float = 7.5,
+                  inversion_prompt: str = "", #needs to be the same as for preprocess
+                  n_fps: int = 10,
+                  progress=gr.Progress(track_tqdm=True)
+):
+    config = {}
+ 
+    config["sd_version"] = "2.1"
+    config["device"] = device
+    config["n_timesteps"] = int(n_timesteps)
+    config["n_frames"] = n_frames
+    config["batch_size"] = batch_size
+    config["guidance_scale"] = gudiance_scale
+    config["prompt"] = prompt
+    config["negative_prompt"] = "ugly, blurry, low res, unrealistic, unaesthetic",
+    config["pnp_attn_t"] = pnp_attn_t
+    config["pnp_f_t"] = pnp_f_t
+    config["pnp_inversion_prompt"] = inversion_prompt
+    
+    
+    if do_inversion:
+        frames, latents, inverted_latents, do_inversion =  preprocess_and_invert(
+                          input_video,
+                          frames,
+                          latents,
+                          inverted_latents,
+                          seed, 
+                          randomize_seed,
+                          do_inversion,
+                          steps,
+                          n_timesteps,
+                          batch_size,
+                          n_frames,
+                          inversion_prompt)
+        do_inversion = False
+        
+    
+    if randomize_seed:
+            seed = randomize_seed_fn()
+    seed_everything(seed)
+    
+    
+    editor = TokenFlow(config=config,pipe=tokenflow_pipe, frames=frames.value, inverted_latents=inverted_latents.value)
+    edited_frames = editor.edit_video()
+
+    save_video(edited_frames, 'tokenflow_PnP_fps_30.mp4', fps=n_fps)
+    # path = export_to_video(edited_frames)
+    return 'tokenflow_PnP_fps_30.mp4', frames, latents, inverted_latents, do_inversion
+
+########
+# demo #
+########
+
+
+intro = """
+<div style="text-align:center">
+<h1 style="font-weight: 1400; text-align: center; margin-bottom: 7px;">
+   TokenFlow - <small>Temporally consistent video editing</small>
+</h1>
+<span>[<a target="_blank" href="https://diffusion-tokenflow.github.io">Project page</a>], [<a target="_blank" href="https://github.com/omerbt/TokenFlow">GitHub</a>], [<a target="_blank" href="https://huggingface.co/papers/2307.10373">Paper</a>]</span>
+<div style="display:flex; justify-content: center;margin-top: 0.5em">Each edit takes ~5 min <a href="https://huggingface.co/weizmannscience/tokenflow?duplicate=true" target="_blank">
+<img style="margin-top: 0em; margin-bottom: 0em; margin-left: 0.5em" src="https://bit.ly/3CWLGkA" alt="Duplicate Space"></a></div>
+</div>
+"""
+
+
+
+with gr.Blocks(css="style.css") as demo:
+    
+    gr.HTML(intro)
+    frames = gr.State()
+    inverted_latents = gr.State()
+    latents = gr.State()
+    do_inversion = gr.State(value=True)
+
+    with gr.Row():
+        input_video = gr.Video(label="Input Video", interactive=True, elem_id="input_video")
+        output_video = gr.Video(label="Edited Video", interactive=False, elem_id="output_video")
+        input_video.style(height=365, width=365)
+        output_video.style(height=365, width=365)
+
+
+    with gr.Row():
+            prompt = gr.Textbox(
+                            label="Describe your edited video",
+                            max_lines=1, value=""
+                        )
+    # with gr.Group(visible=False) as share_btn_container:
+        # with gr.Group(elem_id="share-btn-container"):
+        #     community_icon = gr.HTML(community_icon_html, visible=True)
+        #     loading_icon = gr.HTML(loading_icon_html, visible=False)
+        #     share_button = gr.Button("Share to community", elem_id="share-btn", visible=True)
+        
+   
+    # with gr.Row():
+    #     inversion_progress = gr.Textbox(visible=False, label="Inversion progress")
+               
+    with gr.Row():
+        run_button = gr.Button("Edit your video!", visible=True)
+
+    with gr.Accordion("Advanced Options", open=False):
+        with gr.Tabs() as tabs:
+            with gr.TabItem('General options'):
+                with gr.Row():
+                    with gr.Column(min_width=100):
+                        seed = gr.Number(value=0, precision=0, label="Seed", interactive=True)
+                        randomize_seed = gr.Checkbox(label='Randomize seed', value=False)
+                        gudiance_scale = gr.Slider(label='Guidance Scale', minimum=1, maximum=30,
+                                              value=7.5, step=0.5, interactive=True)
+                        steps = gr.Slider(label='Inversion steps', minimum=10, maximum=500,
+                                              value=500, step=1, interactive=True)
+                        
+                    with gr.Column(min_width=100):
+                        inversion_prompt = gr.Textbox(lines=1, label="Inversion prompt", interactive=True, placeholder="")
+                        batch_size = gr.Slider(label='Batch size', minimum=1, maximum=10,
+                                              value=8, step=1, interactive=True)
+                        n_frames = gr.Slider(label='Num frames', minimum=2, maximum=200,
+                                              value=24, step=1, interactive=True)
+                        n_timesteps = gr.Slider(label='Diffusion steps', minimum=25, maximum=100,
+                                              value=50, step=25, interactive=True)
+                        n_fps = gr.Slider(label='Frames per second', minimum=1, maximum=60,
+                                              value=10, step=1, interactive=True)
+                        
+            with gr.TabItem('Plug-and-Play Parameters'):
+                 with gr.Column(min_width=100):
+                    pnp_attn_t = gr.Slider(label='pnp attention threshold', minimum=0, maximum=1,
+                                              value=0.5, step=0.5, interactive=True)
+                    pnp_f_t = gr.Slider(label='pnp feature threshold', minimum=0, maximum=1,
+                                              value=0.8, step=0.05, interactive=True)
+                    
+                    
+    input_video.change(
+        fn = reset_do_inversion,
+        outputs = [do_inversion],
+        queue = False)
+
+    inversion_prompt.change(
+        fn = reset_do_inversion,
+        outputs = [do_inversion],
+        queue = False)
+
+    randomize_seed.change(
+        fn = reset_do_inversion,
+        outputs = [do_inversion],
+        queue = False)
+
+    seed.change(
+        fn = reset_do_inversion,
+        outputs = [do_inversion],
+        queue = False)
+
+    
+
+    input_video.upload(
+        fn = reset_do_inversion,
+        outputs = [do_inversion],
+        queue = False).then(fn = preprocess_and_invert,
+          inputs = [input_video,
+                      frames,
+                      latents,
+                      inverted_latents,
+                      seed, 
+                      randomize_seed,
+                      do_inversion,
+                      steps,
+                      n_timesteps,
+                      batch_size,
+                      n_frames,
+                      inversion_prompt
+          ],
+          outputs = [frames,
+                     latents,
+                     inverted_latents,
+                     do_inversion
+              
+          ])
+    
+    run_button.click(fn = edit_with_pnp,
+                     inputs = [input_video,
+                               frames, 
+                              latents,
+                              inverted_latents,
+                              seed,
+                              randomize_seed,
+                              do_inversion,
+                              steps,
+                              prompt,                             
+                              pnp_attn_t,
+                              pnp_f_t,
+                              batch_size,
+                              n_frames,
+                              n_timesteps,
+                              gudiance_scale,
+                              inversion_prompt,
+                              n_fps ],
+                                 outputs = [output_video, frames, latents, inverted_latents, do_inversion]
+                                )
+
+    gr.Examples(
+        examples=get_example(),
+        label='Examples',
+        inputs=[input_video],
+        outputs=[output_video]
+    )
+
+demo.queue()
+demo.launch()

examples/running_dog.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:904a4f165d70dd46164e27f3b279869debcfb01797aa4c3f0f3ae6fabea8631d
+size 1443849

examples/woman-running.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c622e40ca8b01a6a678eae719be376cb7969a2b7460955706a214720104880b7
+size 1285788

preprocess_utils.py

@@ -0,0 +1,375 @@
+from transformers import CLIPTextModel, CLIPTokenizer, logging
+from diffusers import AutoencoderKL, UNet2DConditionModel, DDIMScheduler
+# suppress partial model loading warning
+logging.set_verbosity_error()
+
+import os
+from tqdm import tqdm, trange
+import torch
+import torch.nn as nn
+import argparse
+from torchvision.io import write_video
+from pathlib import Path
+from utils import *
+import torchvision.transforms as T
+
+
+def get_timesteps(scheduler, num_inference_steps, strength, device):
+    # get the original timestep using init_timestep
+    init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
+
+    t_start = max(num_inference_steps - init_timestep, 0)
+    timesteps = scheduler.timesteps[t_start:]
+
+    return timesteps, num_inference_steps - t_start
+
+
+class Preprocess(nn.Module):
+    def __init__(self, device, opt, vae, tokenizer, text_encoder, unet,scheduler, hf_key=None):
+        super().__init__()
+
+        self.device = device
+        self.sd_version = opt["sd_version"]
+        self.use_depth = False
+        self.config = opt
+
+        print(f'[INFO] loading stable diffusion...')
+        if hf_key is not None:
+            print(f'[INFO] using hugging face custom model key: {hf_key}')
+            model_key = hf_key
+        elif self.sd_version == '2.1':
+            model_key = "stabilityai/stable-diffusion-2-1-base"
+        elif self.sd_version == '2.0':
+            model_key = "stabilityai/stable-diffusion-2-base"
+        elif self.sd_version == '1.5' or self.sd_version == 'ControlNet':
+            model_key = "runwayml/stable-diffusion-v1-5"
+        elif self.sd_version == 'depth':
+            model_key = "stabilityai/stable-diffusion-2-depth"
+        else:
+            raise ValueError(f'Stable-diffusion version {self.sd_version} not supported.')
+        
+        self.model_key = model_key
+        
+        # Create model
+        # self.vae = AutoencoderKL.from_pretrained(model_key, subfolder="vae", revision="fp16",
+        #                                          torch_dtype=torch.float16).to(self.device)
+        # self.tokenizer = CLIPTokenizer.from_pretrained(model_key, subfolder="tokenizer")
+        # self.text_encoder = CLIPTextModel.from_pretrained(model_key, subfolder="text_encoder", revision="fp16",
+        #                                                   torch_dtype=torch.float16).to(self.device)
+        # self.unet = UNet2DConditionModel.from_pretrained(model_key, subfolder="unet", revision="fp16",
+        #                                            torch_dtype=torch.float16).to(self.device)
+        
+        self.vae = vae
+        self.tokenizer = tokenizer
+        self.text_encoder = text_encoder
+        self.unet = unet
+        self.scheduler=scheduler
+        self.total_inverted_latents = {}
+        
+        self.paths, self.frames, self.latents = self.get_data(self.config["data_path"], self.config["n_frames"])
+        print("self.frames", self.frames.shape)
+        print("self.latents", self.latents.shape)
+        
+        
+        if self.sd_version == 'ControlNet':
+            from diffusers import ControlNetModel, StableDiffusionControlNetPipeline
+            controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16).to(self.device)
+            control_pipe = StableDiffusionControlNetPipeline.from_pretrained(
+                "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
+            ).to(self.device)
+            self.unet = control_pipe.unet
+            self.controlnet = control_pipe.controlnet
+            self.canny_cond = self.get_canny_cond()
+        elif self.sd_version == 'depth':
+            self.depth_maps = self.prepare_depth_maps()
+        self.scheduler = scheduler
+        
+        self.unet.enable_xformers_memory_efficient_attention()
+        print(f'[INFO] loaded stable diffusion!')
+    
+    
+    @torch.no_grad()   
+    def prepare_depth_maps(self, model_type='DPT_Large', device='cuda'):
+        depth_maps = []
+        midas = torch.hub.load("intel-isl/MiDaS", model_type)
+        midas.to(device)
+        midas.eval()
+
+        midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")
+
+        if model_type == "DPT_Large" or model_type == "DPT_Hybrid":
+            transform = midas_transforms.dpt_transform
+        else:
+            transform = midas_transforms.small_transform
+
+        for i in range(len(self.paths)):
+            img = cv2.imread(self.paths[i])
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+            latent_h = img.shape[0] // 8
+            latent_w = img.shape[1] // 8
+            
+            input_batch = transform(img).to(device)
+            prediction = midas(input_batch)
+
+            depth_map = torch.nn.functional.interpolate(
+                prediction.unsqueeze(1),
+                size=(latent_h, latent_w),
+                mode="bicubic",
+                align_corners=False,
+            )
+            depth_min = torch.amin(depth_map, dim=[1, 2, 3], keepdim=True)
+            depth_max = torch.amax(depth_map, dim=[1, 2, 3], keepdim=True)
+            depth_map = 2.0 * (depth_map - depth_min) / (depth_max - depth_min) - 1.0
+            depth_maps.append(depth_map)
+
+        return torch.cat(depth_maps).to(self.device).to(torch.float16)
+    
+    @torch.no_grad()
+    def get_canny_cond(self):
+        canny_cond = []
+        for image in self.frames.cpu().permute(0, 2, 3, 1):
+            image = np.uint8(np.array(255 * image))
+            low_threshold = 100
+            high_threshold = 200
+
+            image = cv2.Canny(image, low_threshold, high_threshold)
+            image = image[:, :, None]
+            image = np.concatenate([image, image, image], axis=2)
+            image = torch.from_numpy((image.astype(np.float32) / 255.0))
+            canny_cond.append(image)
+        canny_cond = torch.stack(canny_cond).permute(0, 3, 1, 2).to(self.device).to(torch.float16)
+        return canny_cond
+    
+    def controlnet_pred(self, latent_model_input, t, text_embed_input, controlnet_cond):
+        down_block_res_samples, mid_block_res_sample = self.controlnet(
+            latent_model_input,
+            t,
+            encoder_hidden_states=text_embed_input,
+            controlnet_cond=controlnet_cond,
+            conditioning_scale=1,
+            return_dict=False,
+        )
+        
+        # apply the denoising network
+        noise_pred = self.unet(
+            latent_model_input,
+            t,
+            encoder_hidden_states=text_embed_input,
+            cross_attention_kwargs={},
+            down_block_additional_residuals=down_block_res_samples,
+            mid_block_additional_residual=mid_block_res_sample,
+            return_dict=False,
+        )[0]
+        return noise_pred
+    
+    @torch.no_grad()
+    def get_text_embeds(self, prompt, negative_prompt, device="cuda"):
+        text_input = self.tokenizer(prompt, padding='max_length', max_length=self.tokenizer.model_max_length,
+                                    truncation=True, return_tensors='pt')
+        text_embeddings = self.text_encoder(text_input.input_ids.to(device))[0]
+        uncond_input = self.tokenizer(negative_prompt, padding='max_length', max_length=self.tokenizer.model_max_length,
+                                      return_tensors='pt')
+        uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(device))[0]
+        text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
+        return text_embeddings
+
+    @torch.no_grad()
+    def decode_latents(self, latents):
+        decoded = []
+        batch_size = 8
+        for b in range(0, latents.shape[0], batch_size):
+                latents_batch = 1 / 0.18215 * latents[b:b + batch_size]
+                imgs = self.vae.decode(latents_batch).sample
+                imgs = (imgs / 2 + 0.5).clamp(0, 1)
+                decoded.append(imgs)
+        return torch.cat(decoded)
+
+    @torch.no_grad()
+    def encode_imgs(self, imgs, batch_size=10, deterministic=True):
+        imgs = 2 * imgs - 1
+        latents = []
+        for i in range(0, len(imgs), batch_size):
+            posterior = self.vae.encode(imgs[i:i + batch_size]).latent_dist
+            latent = posterior.mean if deterministic else posterior.sample()
+            latents.append(latent * 0.18215)
+        latents = torch.cat(latents)
+        return latents
+
+    def get_data(self, frames_path, n_frames):
+        
+        # load frames
+        if not self.config["frames"]:
+            paths =  [f"{frames_path}/%05d.png" % i for i in range(n_frames)]
+            print(paths)
+            if not os.path.exists(paths[0]):
+                paths = [f"{frames_path}/%05d.jpg" % i for i in range(n_frames)]
+            self.paths = paths
+            frames = [Image.open(path).convert('RGB') for path in paths]
+            if frames[0].size[0] == frames[0].size[1]:
+                frames = [frame.resize((512, 512), resample=Image.Resampling.LANCZOS) for frame in frames]
+        else:
+            frames = self.config["frames"][:n_frames]
+        frames = torch.stack([T.ToTensor()(frame) for frame in frames]).to(torch.float16).to(self.device)
+        # encode to latents
+        latents = self.encode_imgs(frames, deterministic=True).to(torch.float16).to(self.device)
+        print("frames", frames.shape)
+        print("latents", latents.shape)
+        
+        if not self.config["frames"]:
+            return paths, frames, latents
+        else:
+            return None, frames, latents
+
+    @torch.no_grad()
+    def ddim_inversion(self, cond, latent_frames, save_path, batch_size, save_latents=True, timesteps_to_save=None):
+        timesteps = reversed(self.scheduler.timesteps)
+        timesteps_to_save = timesteps_to_save if timesteps_to_save is not None else timesteps
+        
+        return_inverted_latents = self.config["frames"] is not None
+        for i, t in enumerate(tqdm(timesteps)):
+            for b in range(0, latent_frames.shape[0], int(batch_size)):
+                x_batch = latent_frames[b:b + batch_size]
+                model_input = x_batch
+                cond_batch = cond.repeat(x_batch.shape[0], 1, 1)
+                if self.sd_version == 'depth':
+                    depth_maps = torch.cat([self.depth_maps[b: b + batch_size]])
+                    model_input = torch.cat([x_batch, depth_maps],dim=1)
+                                                                    
+                alpha_prod_t = self.scheduler.alphas_cumprod[t]
+                alpha_prod_t_prev = (
+                    self.scheduler.alphas_cumprod[timesteps[i - 1]]
+                    if i > 0 else self.scheduler.final_alpha_cumprod
+                )
+
+                mu = alpha_prod_t ** 0.5
+                mu_prev = alpha_prod_t_prev ** 0.5
+                sigma = (1 - alpha_prod_t) ** 0.5
+                sigma_prev = (1 - alpha_prod_t_prev) ** 0.5
+
+                eps = self.unet(model_input, t, encoder_hidden_states=cond_batch).sample if self.sd_version != 'ControlNet' \
+                    else self.controlnet_pred(x_batch, t, cond_batch, torch.cat([self.canny_cond[b: b + batch_size]]))
+                pred_x0 = (x_batch - sigma_prev * eps) / mu_prev
+                latent_frames[b:b + batch_size] = mu * pred_x0 + sigma * eps
+            
+            if return_inverted_latents and t in timesteps_to_save:
+                self.total_inverted_latents[f'noisy_latents_{t}'] = latent_frames.clone()
+    
+            if save_latents and t in timesteps_to_save:
+                torch.save(latent_frames, os.path.join(save_path, 'latents', f'noisy_latents_{t}.pt'))
+        
+        if save_latents:
+            torch.save(latent_frames, os.path.join(save_path, 'latents', f'noisy_latents_{t}.pt'))
+        if return_inverted_latents:
+            self.total_inverted_latents[f'noisy_latents_{t}'] = latent_frames.clone()
+
+        return latent_frames
+
+    @torch.no_grad()
+    def ddim_sample(self, x, cond, batch_size):
+        timesteps = self.scheduler.timesteps
+        for i, t in enumerate(tqdm(timesteps)):
+            for b in range(0, x.shape[0], batch_size):
+                x_batch = x[b:b + batch_size]
+                model_input = x_batch
+                cond_batch = cond.repeat(x_batch.shape[0], 1, 1)
+                
+                if self.sd_version == 'depth':
+                    depth_maps = torch.cat([self.depth_maps[b: b + batch_size]])
+                    model_input = torch.cat([x_batch, depth_maps],dim=1)
+                
+                alpha_prod_t = self.scheduler.alphas_cumprod[t]
+                alpha_prod_t_prev = (
+                    self.scheduler.alphas_cumprod[timesteps[i + 1]]
+                    if i < len(timesteps) - 1
+                    else self.scheduler.final_alpha_cumprod
+                )
+                mu = alpha_prod_t ** 0.5
+                sigma = (1 - alpha_prod_t) ** 0.5
+                mu_prev = alpha_prod_t_prev ** 0.5
+                sigma_prev = (1 - alpha_prod_t_prev) ** 0.5
+
+                eps = self.unet(model_input, t, encoder_hidden_states=cond_batch).sample if self.sd_version != 'ControlNet' \
+                    else self.controlnet_pred(x_batch, t, cond_batch, torch.cat([self.canny_cond[b: b + batch_size]]))
+
+                pred_x0 = (x_batch - sigma * eps) / mu
+                x[b:b + batch_size] = mu_prev * pred_x0 + sigma_prev * eps
+        return x
+
+    @torch.no_grad()
+    def extract_latents(self, 
+                        num_steps,
+                        save_path,
+                        batch_size,
+                        timesteps_to_save,
+                        inversion_prompt='',
+                       reconstruct=False):
+        self.scheduler.set_timesteps(num_steps)
+        cond = self.get_text_embeds(inversion_prompt, "")[1].unsqueeze(0)
+        latent_frames = self.latents
+        print("latent_frames", latent_frames.shape)
+  
+        inverted_x= self.ddim_inversion(cond,
+                                         latent_frames,
+                                         save_path,
+                                         batch_size=batch_size,
+                                         save_latents=True if save_path else False,
+                                         timesteps_to_save=timesteps_to_save)
+
+
+       
+        # print("total_inverted_latents", len(total_inverted_latents.keys()))
+        
+        if reconstruct:
+            latent_reconstruction = self.ddim_sample(inverted_x, cond, batch_size=batch_size)
+
+            rgb_reconstruction = self.decode_latents(latent_reconstruction)
+            return self.frames, self.latents, self.total_inverted_latents, rgb_reconstruction
+
+        return self.frames, self.latents, self.total_inverted_latents, None
+
+
+def prep(opt):
+    # timesteps to save
+    if opt["sd_version"] == '2.1':
+        model_key = "stabilityai/stable-diffusion-2-1-base"
+    elif opt["sd_version"] == '2.0':
+        model_key = "stabilityai/stable-diffusion-2-base"
+    elif opt["sd_version"] == '1.5' or opt["sd_version"] == 'ControlNet':
+        model_key = "runwayml/stable-diffusion-v1-5"
+    elif opt["sd_version"] == 'depth':
+        model_key = "stabilityai/stable-diffusion-2-depth"
+    toy_scheduler = DDIMScheduler.from_pretrained(model_key, subfolder="scheduler")
+    toy_scheduler.set_timesteps(opt["save_steps"])
+    timesteps_to_save, num_inference_steps = get_timesteps(toy_scheduler, num_inference_steps=opt["save_steps"],
+                                                           strength=1.0,
+                                                           device=device)
+
+    seed_everything(opt["seed"])
+    if not opt["frames"]: # original non demo setting
+        save_path = os.path.join(opt["save_dir"],
+                                 f'sd_{opt["sd_version"]}',
+                                 Path(opt["data_path"]).stem,
+                                 f'steps_{opt["steps"]}',
+                                 f'nframes_{opt["n_frames"]}') 
+        os.makedirs(os.path.join(save_path, f'latents'), exist_ok=True)
+        add_dict_to_yaml_file(os.path.join(opt["save_dir"], 'inversion_prompts.yaml'), Path(opt["data_path"]).stem, opt["inversion_prompt"])    
+        # save inversion prompt in a txt file
+        with open(os.path.join(save_path, 'inversion_prompt.txt'), 'w') as f:
+            f.write(opt["inversion_prompt"])
+    else:
+        save_path = None
+    
+    model = Preprocess(device, opt)
+  
+    frames, latents, total_inverted_latents, rgb_reconstruction = model.extract_latents(
+                                         num_steps=model.config["steps"],
+                                         save_path=save_path,
+                                         batch_size=model.config["batch_size"],
+                                         timesteps_to_save=timesteps_to_save,
+                                         inversion_prompt=model.config["inversion_prompt"],
+    )
+
+    
+    return frames, latents, total_inverted_latents, rgb_reconstruction
+

requirements.txt

@@ -0,0 +1,14 @@
+pillow
+diffusers
+ftfy
+transformers
+opencv-python 
+tqdm
+numpy
+pyyaml
+accelerate
+xformers
+tensorboard
+kornia
+av
+torchvision==0.15.2

style.css

@@ -0,0 +1,91 @@
+/*
+This CSS file is modified from:
+https://huggingface.co/spaces/DeepFloyd/IF/blob/main/style.css
+*/
+
+h1 {
+  text-align: center;
+}
+
+.gradio-container {
+  font-family: 'IBM Plex Sans', sans-serif;
+}
+
+.gr-button {
+  color: white;
+  border-color: black;
+  background: black;
+}
+
+input[type='range'] {
+  accent-color: black;
+}
+
+.dark input[type='range'] {
+  accent-color: #dfdfdf;
+}
+
+.container {
+  max-width: 730px;
+  margin: auto;
+}
+
+.gr-button:focus {
+  border-color: rgb(147 197 253 / var(--tw-border-opacity));
+  outline: none;
+  box-shadow: var(--tw-ring-offset-shadow), var(--tw-ring-shadow), var(--tw-shadow, 0 0 #0000);
+  --tw-border-opacity: 1;
+  --tw-ring-offset-shadow: var(--tw-ring-inset) 0 0 0 var(--tw-ring-offset-width) var(--tw-ring-offset-color);
+  --tw-ring-shadow: var(--tw-ring-inset) 0 0 0 calc(3px var(--tw-ring-offset-width)) var(--tw-ring-color);
+  --tw-ring-color: rgb(191 219 254 / var(--tw-ring-opacity));
+  --tw-ring-opacity: .5;
+}
+
+.gr-form {
+  flex: 1 1 50%;
+  border-top-right-radius: 0;
+  border-bottom-right-radius: 0;
+}
+
+#prompt-container {
+  gap: 0;
+}
+
+#prompt-text-input,
+#negative-prompt-text-input {
+  padding: .45rem 0.625rem
+}
+
+/* #component-16 {
+  border-top-width: 1px !important;
+  margin-top: 1em
+} */
+
+.image_duplication {
+  position: absolute;
+  width: 100px;
+  left: 50px
+}
+
+#component-0 {
+  max-width: 730px;
+  margin: auto;
+  padding-top: 1.5rem;
+}
+
+#share-btn-container {
+    display: flex; padding-left: 0.5rem !important; padding-right: 0.5rem !important; background-color: #000000; justify-content: center; align-items: center; border-radius: 9999px !important; width: 13rem; margin-left: auto;
+}
+#share-btn {
+    all: initial; color: #ffffff;font-weight: 600; cursor:pointer; font-family: 'IBM Plex Sans', sans-serif; margin-left: 0.5rem !important; padding-top: 0.25rem !important; padding-bottom: 0.25rem !important;
+}
+#share-btn * {
+    all: unset;
+}
+#share-btn-container div:nth-child(-n+2){
+    width: auto !important;
+    min-height: 0px !important;
+}
+#share-btn-container .wrap {
+    display: none !important;
+}

tokenflow_pnp.py

@@ -0,0 +1,364 @@
+import glob
+import os
+import numpy as np
+import cv2
+from pathlib import Path
+import torch
+import torch.nn as nn
+import torchvision.transforms as T
+import argparse
+from PIL import Image
+import yaml
+from tqdm import tqdm
+from transformers import logging
+from diffusers import DDIMScheduler, StableDiffusionPipeline
+
+from tokenflow_utils import *
+from utils import save_video, seed_everything
+
+# suppress partial model loading warning
+logging.set_verbosity_error()
+
+VAE_BATCH_SIZE = 10
+
+
+class TokenFlow(nn.Module):
+    def __init__(self, config, 
+                 pipe,
+                 frames=None,
+                # latents = None,
+                inverted_latents = None):
+        super().__init__()
+        self.config = config
+        self.device = config["device"]
+        
+        sd_version = config["sd_version"]
+        self.sd_version = sd_version
+        if sd_version == '2.1':
+            model_key = "stabilityai/stable-diffusion-2-1-base"
+        elif sd_version == '2.0':
+            model_key = "stabilityai/stable-diffusion-2-base"
+        elif sd_version == '1.5':
+            model_key = "runwayml/stable-diffusion-v1-5"
+        elif sd_version == 'depth':
+            model_key = "stabilityai/stable-diffusion-2-depth"
+        else:
+            raise ValueError(f'Stable-diffusion version {sd_version} not supported.')
+
+        # Create SD models
+        print('Loading SD model')
+
+        # pipe = StableDiffusionPipeline.from_pretrained(model_key, torch_dtype=torch.float16).to("cuda")
+        # pipe.enable_xformers_memory_efficient_attention()
+
+        self.vae = pipe.vae
+        self.tokenizer = pipe.tokenizer
+        self.text_encoder = pipe.text_encoder
+        self.unet = pipe.unet
+
+        self.scheduler = DDIMScheduler.from_pretrained(model_key, subfolder="scheduler")
+        self.scheduler.set_timesteps(config["n_timesteps"], device=self.device)
+        print('SD model loaded')
+        
+        # data
+        self.frames, self.inverted_latents  = frames, inverted_latents
+        self.latents_path = self.get_latents_path()
+        
+        # load frames
+        self.paths, self.frames, self.latents, self.eps = self.get_data()
+        
+        if self.sd_version == 'depth':
+            self.depth_maps = self.prepare_depth_maps()
+
+        self.text_embeds = self.get_text_embeds(config["prompt"], config["negative_prompt"])
+        # pnp_inversion_prompt = self.get_pnp_inversion_prompt()
+        self.pnp_guidance_embeds = self.get_text_embeds(config["pnp_inversion_prompt"], config["pnp_inversion_prompt"]).chunk(2)[0]
+    
+    @torch.no_grad()   
+    def prepare_depth_maps(self, model_type='DPT_Large', device='cuda'):
+        depth_maps = []
+        midas = torch.hub.load("intel-isl/MiDaS", model_type)
+        midas.to(device)
+        midas.eval()
+
+        midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")
+
+        if model_type == "DPT_Large" or model_type == "DPT_Hybrid":
+            transform = midas_transforms.dpt_transform
+        else:
+            transform = midas_transforms.small_transform
+
+        for i in range(len(self.paths)):
+            img = cv2.imread(self.paths[i])
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+            
+            latent_h = img.shape[0] // 8
+            latent_w = img.shape[1] // 8
+            
+            input_batch = transform(img).to(device)
+            prediction = midas(input_batch)
+
+            depth_map = torch.nn.functional.interpolate(
+                prediction.unsqueeze(1),
+                size=(latent_h, latent_w),
+                mode="bicubic",
+                align_corners=False,
+            )
+            depth_min = torch.amin(depth_map, dim=[1, 2, 3], keepdim=True)
+            depth_max = torch.amax(depth_map, dim=[1, 2, 3], keepdim=True)
+            depth_map = 2.0 * (depth_map - depth_min) / (depth_max - depth_min) - 1.0
+            depth_maps.append(depth_map)
+
+        return torch.cat(depth_maps).to(torch.float16).to(self.device)
+    
+    def get_pnp_inversion_prompt(self):
+        inv_prompts_path = os.path.join(str(Path(self.latents_path).parent), 'inversion_prompt.txt')
+        # read inversion prompt
+        with open(inv_prompts_path, 'r') as f:
+            inv_prompt = f.read()
+        return inv_prompt
+
+    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:
+            n_frames = self.frames.shape[0]
+            self.config["n_frames"] = min(n_frames, self.config["n_frames"])
+        
+        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
+
+    @torch.no_grad()
+    def get_text_embeds(self, prompt, negative_prompt, batch_size=1):
+        # Tokenize text and get embeddings
+        text_input = self.tokenizer(prompt, padding='max_length', max_length=self.tokenizer.model_max_length,
+                                    truncation=True, return_tensors='pt')
+        text_embeddings = self.text_encoder(text_input.input_ids.to(self.device))[0]
+
+        # Do the same for unconditional embeddings
+        uncond_input = self.tokenizer(negative_prompt, padding='max_length', max_length=self.tokenizer.model_max_length,
+                                      return_tensors='pt')
+
+        uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(self.device))[0]
+
+        # Cat for final embeddings
+        text_embeddings = torch.cat([uncond_embeddings] * batch_size + [text_embeddings] * batch_size)
+        return text_embeddings
+
+    @torch.no_grad()
+    def encode_imgs(self, imgs, batch_size=VAE_BATCH_SIZE, deterministic=False):
+        imgs = 2 * imgs - 1
+        latents = []
+        for i in range(0, len(imgs), batch_size):
+            posterior = self.vae.encode(imgs[i:i + batch_size]).latent_dist
+            latent = posterior.mean if deterministic else posterior.sample()
+            latents.append(latent * 0.18215)
+        latents = torch.cat(latents)
+        return latents
+
+    @torch.no_grad()
+    def decode_latents(self, latents, batch_size=VAE_BATCH_SIZE):
+        latents = 1 / 0.18215 * latents
+        imgs = []
+        for i in range(0, len(latents), batch_size):
+            imgs.append(self.vae.decode(latents[i:i + batch_size]).sample)
+        imgs = torch.cat(imgs)
+        imgs = (imgs / 2 + 0.5).clamp(0, 1)
+        return imgs
+
+    
+    def get_data(self):
+        read_from_files = self.frames is None
+        # read_from_files = True
+        if read_from_files:
+            # load frames
+            paths = [os.path.join(self.config["data_path"], "%05d.jpg" % idx) for idx in
+                                   range(self.config["n_frames"])]
+            if not os.path.exists(paths[0]):
+                paths = [os.path.join(self.config["data_path"], "%05d.png" % idx) for idx in
+                                       range(self.config["n_frames"])]
+            frames = [Image.open(paths[idx]).convert('RGB') for idx in range(self.config["n_frames"])]
+            if frames[0].size[0] == frames[0].size[1]:
+                frames = [frame.resize((512, 512), resample=Image.Resampling.LANCZOS) for frame in frames]
+            frames = torch.stack([T.ToTensor()(frame) for frame in frames]).to(torch.float16).to(self.device)
+            save_video(frames, f'{self.config["output_path"]}/input_fps10.mp4', fps=10)
+            save_video(frames, f'{self.config["output_path"]}/input_fps20.mp4', fps=20)
+            save_video(frames, f'{self.config["output_path"]}/input_fps30.mp4', fps=30)
+        else:
+            frames = self.frames
+        # 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 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
+
+    @torch.no_grad()
+    def denoise_step(self, x, t, indices):
+        # register the time step and features in pnp injection modules
+        read_files = self.inverted_latents is None
+
+        if read_files:
+            source_latents = load_source_latents_t(t, self.latents_path)[indices]
+
+        else:
+            source_latents = self.inverted_latents[f'noisy_latents_{t}'][indices]
+
+        latent_model_input = torch.cat([source_latents] + ([x] * 2))
+        if self.sd_version == 'depth':
+            latent_model_input = torch.cat([latent_model_input, torch.cat([self.depth_maps[indices]] * 3)], dim=1)
+
+        register_time(self, t.item())
+
+        # compute text embeddings
+        text_embed_input = torch.cat([self.pnp_guidance_embeds.repeat(len(indices), 1, 1),
+                                      torch.repeat_interleave(self.text_embeds, len(indices), dim=0)])
+
+        # apply the denoising network
+        noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embed_input)['sample']
+
+        # perform guidance
+        _, noise_pred_uncond, noise_pred_cond = noise_pred.chunk(3)
+        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']
+        return denoised_latent
+    
+    @torch.autocast(dtype=torch.float16, device_type='cuda')
+    def batched_denoise_step(self, x, t, indices):
+        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) 
+            
+        register_pivotal(self, True)
+        self.denoise_step(x[pivotal_idx], t, indices[pivotal_idx])
+        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]))
+        denoised_latents = torch.cat(denoised_latents)
+        return denoised_latents
+
+    def init_method(self, conv_injection_t, qk_injection_t):
+        self.qk_injection_timesteps = self.scheduler.timesteps[:qk_injection_t] if qk_injection_t >= 0 else []
+        self.conv_injection_timesteps = self.scheduler.timesteps[:conv_injection_t] if conv_injection_t >= 0 else []
+        register_extended_attention_pnp(self, self.qk_injection_timesteps)
+        register_conv_injection(self, self.conv_injection_timesteps)
+        set_tokenflow(self.unet)
+
+    def save_vae_recon(self):
+        os.makedirs(f'{self.config["output_path"]}/vae_recon', exist_ok=True)
+        decoded = self.decode_latents(self.latents)
+        for i in range(len(decoded)):
+            T.ToPILImage()(decoded[i]).save(f'{self.config["output_path"]}/vae_recon/%05d.png' % i)
+        save_video(decoded, f'{self.config["output_path"]}/vae_recon_10.mp4', fps=10)
+        save_video(decoded, f'{self.config["output_path"]}/vae_recon_20.mp4', fps=20)
+        save_video(decoded, f'{self.config["output_path"]}/vae_recon_30.mp4', fps=30)
+
+    def edit_video(self):
+        save_files = self.inverted_latents is None # if we're in the original non-demo setting
+        if save_files: 
+            os.makedirs(f'{self.config["output_path"]}/img_ode', exist_ok=True)
+            self.save_vae_recon()
+        # self.save_vae_recon()
+        pnp_f_t = int(self.config["n_timesteps"] * self.config["pnp_f_t"])
+        pnp_attn_t = int(self.config["n_timesteps"] * self.config["pnp_attn_t"])
+        
+        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])
+        edited_frames = self.sample_loop(noisy_latents, torch.arange(self.config["n_frames"]))
+        
+        if save_files:
+            save_video(edited_frames, f'{self.config["output_path"]}/tokenflow_PnP_fps_10.mp4')
+            save_video(edited_frames, f'{self.config["output_path"]}/tokenflow_PnP_fps_20.mp4', fps=20)
+            save_video(edited_frames, f'{self.config["output_path"]}/tokenflow_PnP_fps_30.mp4', fps=30)
+            print('Done!')
+        else:
+            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
+        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)
+        
+        decoded_latents = self.decode_latents(x)
+        if save_files:
+            for i in range(len(decoded_latents)):
+                T.ToPILImage()(decoded_latents[i]).save(f'{self.config["output_path"]}/img_ode/%05d.png' % i)
+
+        return decoded_latents
+
+
+# def run(config):
+#     seed_everything(config["seed"])
+#     print(config)
+#     editor = TokenFlow(config)
+#     editor.edit_video()
+
+
+# if __name__ == '__main__':
+#     parser = argparse.ArgumentParser()
+#     parser.add_argument('--config_path', type=str, default='configs/config_pnp.yaml')
+#     opt = parser.parse_args()
+#     with open(opt.config_path, "r") as f:
+#         config = yaml.safe_load(f)
+#     config["output_path"] = os.path.join(config["output_path"] + f'_pnp_SD_{config["sd_version"]}',
+#                                              Path(config["data_path"]).stem,
+#                                              config["prompt"][:240],
+#                                              f'attn_{config["pnp_attn_t"]}_f_{config["pnp_f_t"]}',
+#                                              f'batch_size_{str(config["batch_size"])}',
+#                                              str(config["n_timesteps"]),
+#     )
+#     os.makedirs(config["output_path"], exist_ok=True)
+#     print(config["data_path"])
+#     assert os.path.exists(config["data_path"]), "Data path does not exist"
+#     with open(os.path.join(config["output_path"], "config.yaml"), "w") as f:
+#         yaml.dump(config, f)
+#     run(config)

tokenflow_utils.py

@@ -0,0 +1,448 @@
+from typing import Type
+import torch
+import os
+
+from utils import isinstance_str, batch_cosine_sim
+
+def register_pivotal(diffusion_model, is_pivotal):
+    for _, module in diffusion_model.named_modules():
+        # If for some reason this has a different name, create an issue and I'll fix it
+        if isinstance_str(module, "BasicTransformerBlock"):
+            setattr(module, "pivotal_pass", is_pivotal)
+            
+def register_batch_idx(diffusion_model, batch_idx):
+    for _, module in diffusion_model.named_modules():
+        # If for some reason this has a different name, create an issue and I'll fix it
+        if isinstance_str(module, "BasicTransformerBlock"):
+            setattr(module, "batch_idx", batch_idx)
+
+
+def register_time(model, t):
+    conv_module = model.unet.up_blocks[1].resnets[1]
+    setattr(conv_module, 't', t)
+    down_res_dict = {0: [0, 1], 1: [0, 1], 2: [0, 1]}
+    up_res_dict = {1: [0, 1, 2], 2: [0, 1, 2], 3: [0, 1, 2]}
+    for res in up_res_dict:
+        for block in up_res_dict[res]:
+            module = model.unet.up_blocks[res].attentions[block].transformer_blocks[0].attn1
+            setattr(module, 't', t)
+            module = model.unet.up_blocks[res].attentions[block].transformer_blocks[0].attn2
+            setattr(module, 't', t)
+    for res in down_res_dict:
+        for block in down_res_dict[res]:
+            module = model.unet.down_blocks[res].attentions[block].transformer_blocks[0].attn1
+            setattr(module, 't', t)
+            module = model.unet.down_blocks[res].attentions[block].transformer_blocks[0].attn2
+            setattr(module, 't', t)
+    module = model.unet.mid_block.attentions[0].transformer_blocks[0].attn1
+    setattr(module, 't', t)
+    module = model.unet.mid_block.attentions[0].transformer_blocks[0].attn2
+    setattr(module, 't', t)
+
+
+def load_source_latents_t(t, latents_path):
+    latents_t_path = os.path.join(latents_path, f'noisy_latents_{t}.pt')
+    assert os.path.exists(latents_t_path), f'Missing latents at t {t} path {latents_t_path}'
+    latents = torch.load(latents_t_path)
+    return latents
+
+def register_conv_injection(model, injection_schedule):
+    def conv_forward(self):
+        def forward(input_tensor, temb):
+            hidden_states = input_tensor
+
+            hidden_states = self.norm1(hidden_states)
+            hidden_states = self.nonlinearity(hidden_states)
+
+            if self.upsample is not None:
+                # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
+                if hidden_states.shape[0] >= 64:
+                    input_tensor = input_tensor.contiguous()
+                    hidden_states = hidden_states.contiguous()
+                input_tensor = self.upsample(input_tensor)
+                hidden_states = self.upsample(hidden_states)
+            elif self.downsample is not None:
+                input_tensor = self.downsample(input_tensor)
+                hidden_states = self.downsample(hidden_states)
+
+            hidden_states = self.conv1(hidden_states)
+
+            if temb is not None:
+                temb = self.time_emb_proj(self.nonlinearity(temb))[:, :, None, None]
+
+            if temb is not None and self.time_embedding_norm == "default":
+                hidden_states = hidden_states + temb
+
+            hidden_states = self.norm2(hidden_states)
+
+            if temb is not None and self.time_embedding_norm == "scale_shift":
+                scale, shift = torch.chunk(temb, 2, dim=1)
+                hidden_states = hidden_states * (1 + scale) + shift
+
+            hidden_states = self.nonlinearity(hidden_states)
+
+            hidden_states = self.dropout(hidden_states)
+            hidden_states = self.conv2(hidden_states)
+            if self.injection_schedule is not None and (self.t in self.injection_schedule or self.t == 1000):
+                source_batch_size = int(hidden_states.shape[0] // 3)
+                # inject unconditional
+                hidden_states[source_batch_size:2 * source_batch_size] = hidden_states[:source_batch_size]
+                # inject conditional
+                hidden_states[2 * source_batch_size:] = hidden_states[:source_batch_size]
+
+            if self.conv_shortcut is not None:
+                input_tensor = self.conv_shortcut(input_tensor)
+
+            output_tensor = (input_tensor + hidden_states) / self.output_scale_factor
+
+            return output_tensor
+
+        return forward
+
+    conv_module = model.unet.up_blocks[1].resnets[1]
+    conv_module.forward = conv_forward(conv_module)
+    setattr(conv_module, 'injection_schedule', injection_schedule)
+
+def register_extended_attention_pnp(model, injection_schedule):
+    def sa_forward(self):
+        to_out = self.to_out
+        if type(to_out) is torch.nn.modules.container.ModuleList:
+            to_out = self.to_out[0]
+        else:
+            to_out = self.to_out
+
+        def forward(x, encoder_hidden_states=None):
+            batch_size, sequence_length, dim = x.shape
+            h = self.heads
+            n_frames = batch_size // 3
+            is_cross = encoder_hidden_states is not None
+            encoder_hidden_states = encoder_hidden_states if is_cross else x
+            q = self.to_q(x)
+            k = self.to_k(encoder_hidden_states)
+            v = self.to_v(encoder_hidden_states)
+
+            if self.injection_schedule is not None and (self.t in self.injection_schedule or self.t == 1000):
+                # inject unconditional
+                q[n_frames:2 * n_frames] = q[:n_frames]
+                k[n_frames:2 * n_frames] = k[:n_frames]
+                # inject conditional
+                q[2 * n_frames:] = q[:n_frames]
+                k[2 * n_frames:] = k[:n_frames]
+
+            k_source = k[:n_frames]
+            k_uncond = k[n_frames:2 * n_frames].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1)
+            k_cond = k[2 * n_frames:].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1)
+
+            v_source = v[:n_frames]
+            v_uncond = v[n_frames:2 * n_frames].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1)
+            v_cond = v[2 * n_frames:].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1)
+
+            q_source = self.head_to_batch_dim(q[:n_frames])
+            q_uncond = self.head_to_batch_dim(q[n_frames:2 * n_frames])
+            q_cond = self.head_to_batch_dim(q[2 * n_frames:])
+            k_source = self.head_to_batch_dim(k_source)
+            k_uncond = self.head_to_batch_dim(k_uncond)
+            k_cond = self.head_to_batch_dim(k_cond)
+            v_source = self.head_to_batch_dim(v_source)
+            v_uncond = self.head_to_batch_dim(v_uncond)
+            v_cond = self.head_to_batch_dim(v_cond)
+
+
+            q_src = q_source.view(n_frames, h, sequence_length, dim // h)
+            k_src = k_source.view(n_frames, h, sequence_length, dim // h)
+            v_src = v_source.view(n_frames, h, sequence_length, dim // h)
+            q_uncond = q_uncond.view(n_frames, h, sequence_length, dim // h)
+            k_uncond = k_uncond.view(n_frames, h, sequence_length * n_frames, dim // h)
+            v_uncond = v_uncond.view(n_frames, h, sequence_length * n_frames, dim // h)
+            q_cond = q_cond.view(n_frames, h, sequence_length, dim // h)
+            k_cond = k_cond.view(n_frames, h, sequence_length * n_frames, dim // h)
+            v_cond = v_cond.view(n_frames, h, sequence_length * n_frames, dim // h)
+
+            out_source_all = []
+            out_uncond_all = []
+            out_cond_all = []
+            
+            single_batch = n_frames<=12
+            b = n_frames if single_batch else 1
+
+            for frame in range(0, n_frames, b):
+                out_source = []
+                out_uncond = []
+                out_cond = []
+                for j in range(h):
+                    sim_source_b = torch.bmm(q_src[frame: frame+ b, j], k_src[frame: frame+ b, j].transpose(-1, -2)) * self.scale
+                    sim_uncond_b = torch.bmm(q_uncond[frame: frame+ b, j], k_uncond[frame: frame+ b, j].transpose(-1, -2)) * self.scale
+                    sim_cond = torch.bmm(q_cond[frame: frame+ b, j], k_cond[frame: frame+ b, j].transpose(-1, -2)) * self.scale
+
+                    out_source.append(torch.bmm(sim_source_b.softmax(dim=-1), v_src[frame: frame+ b, j]))
+                    out_uncond.append(torch.bmm(sim_uncond_b.softmax(dim=-1), v_uncond[frame: frame+ b, j]))
+                    out_cond.append(torch.bmm(sim_cond.softmax(dim=-1), v_cond[frame: frame+ b, j]))
+
+                out_source = torch.cat(out_source, dim=0)
+                out_uncond = torch.cat(out_uncond, dim=0) 
+                out_cond = torch.cat(out_cond, dim=0) 
+                if single_batch:
+                    out_source = out_source.view(h, n_frames,sequence_length, dim // h).permute(1, 0, 2, 3).reshape(h * n_frames, sequence_length, -1)
+                    out_uncond = out_uncond.view(h, n_frames,sequence_length, dim // h).permute(1, 0, 2, 3).reshape(h * n_frames, sequence_length, -1)
+                    out_cond = out_cond.view(h, n_frames,sequence_length, dim // h).permute(1, 0, 2, 3).reshape(h * n_frames, sequence_length, -1)
+                out_source_all.append(out_source)
+                out_uncond_all.append(out_uncond)
+                out_cond_all.append(out_cond)
+            
+            out_source = torch.cat(out_source_all, dim=0)
+            out_uncond = torch.cat(out_uncond_all, dim=0)
+            out_cond = torch.cat(out_cond_all, dim=0)
+                
+            out = torch.cat([out_source, out_uncond, out_cond], dim=0)
+            out = self.batch_to_head_dim(out)
+
+            return to_out(out)
+
+        return forward
+
+    for _, module in model.unet.named_modules():
+        if isinstance_str(module, "BasicTransformerBlock"):
+            module.attn1.forward = sa_forward(module.attn1)
+            setattr(module.attn1, 'injection_schedule', [])
+
+    res_dict = {1: [1, 2], 2: [0, 1, 2], 3: [0, 1, 2]}
+    # we are injecting attention in blocks 4 - 11 of the decoder, so not in the first block of the lowest resolution
+    for res in res_dict:
+        for block in res_dict[res]:
+            module = model.unet.up_blocks[res].attentions[block].transformer_blocks[0].attn1
+            module.forward = sa_forward(module)
+            setattr(module, 'injection_schedule', injection_schedule)
+
+def register_extended_attention(model):
+    def sa_forward(self):
+        to_out = self.to_out
+        if type(to_out) is torch.nn.modules.container.ModuleList:
+            to_out = self.to_out[0]
+        else:
+            to_out = self.to_out
+
+        def forward(x, encoder_hidden_states=None):
+            batch_size, sequence_length, dim = x.shape
+            h = self.heads
+            n_frames = batch_size // 3
+            is_cross = encoder_hidden_states is not None
+            encoder_hidden_states = encoder_hidden_states if is_cross else x
+            q = self.to_q(x)
+            k = self.to_k(encoder_hidden_states)
+            v = self.to_v(encoder_hidden_states)
+
+            k_source = k[:n_frames]
+            k_uncond = k[n_frames: 2*n_frames].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1)
+            k_cond = k[2*n_frames:].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1)
+            v_source = v[:n_frames]
+            v_uncond = v[n_frames:2*n_frames].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1)
+            v_cond = v[2*n_frames:].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1)
+
+            q_source = self.head_to_batch_dim(q[:n_frames])
+            q_uncond = self.head_to_batch_dim(q[n_frames: 2*n_frames])
+            q_cond = self.head_to_batch_dim(q[2 * n_frames:])
+            k_source = self.head_to_batch_dim(k_source)
+            k_uncond = self.head_to_batch_dim(k_uncond)
+            k_cond = self.head_to_batch_dim(k_cond)
+            v_source = self.head_to_batch_dim(v_source)
+            v_uncond = self.head_to_batch_dim(v_uncond)
+            v_cond = self.head_to_batch_dim(v_cond)
+
+            out_source = []
+            out_uncond = []
+            out_cond = []
+
+            q_src = q_source.view(n_frames, h, sequence_length, dim // h)
+            k_src = k_source.view(n_frames, h, sequence_length, dim // h)
+            v_src = v_source.view(n_frames, h, sequence_length, dim // h)
+            q_uncond = q_uncond.view(n_frames, h, sequence_length, dim // h)
+            k_uncond = k_uncond.view(n_frames, h, sequence_length * n_frames, dim // h)
+            v_uncond = v_uncond.view(n_frames, h, sequence_length * n_frames, dim // h)
+            q_cond = q_cond.view(n_frames, h, sequence_length, dim // h)
+            k_cond = k_cond.view(n_frames, h, sequence_length * n_frames, dim // h)
+            v_cond = v_cond.view(n_frames, h, sequence_length * n_frames, dim // h)
+
+            for j in range(h):
+                sim_source_b = torch.bmm(q_src[:, j], k_src[:, j].transpose(-1, -2)) * self.scale
+                sim_uncond_b = torch.bmm(q_uncond[:, j], k_uncond[:, j].transpose(-1, -2)) * self.scale
+                sim_cond = torch.bmm(q_cond[:, j], k_cond[:, j].transpose(-1, -2)) * self.scale
+
+                out_source.append(torch.bmm(sim_source_b.softmax(dim=-1), v_src[:, j]))
+                out_uncond.append(torch.bmm(sim_uncond_b.softmax(dim=-1), v_uncond[:, j]))
+                out_cond.append(torch.bmm(sim_cond.softmax(dim=-1), v_cond[:, j]))
+
+            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)
+            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)
+            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)
+
+            out = torch.cat([out_source, out_uncond, out_cond], dim=0)
+            out = self.batch_to_head_dim(out)
+
+            return to_out(out)
+
+        return forward
+
+    for _, module in model.unet.named_modules():
+        if isinstance_str(module, "BasicTransformerBlock"):
+            module.attn1.forward = sa_forward(module.attn1)
+
+    res_dict = {1: [1, 2], 2: [0, 1, 2], 3: [0, 1, 2]}
+    # we are injecting attention in blocks 4 - 11 of the decoder, so not in the first block of the lowest resolution
+    for res in res_dict:
+        for block in res_dict[res]:
+            module = model.unet.up_blocks[res].attentions[block].transformer_blocks[0].attn1
+            module.forward = sa_forward(module)
+
+def make_tokenflow_attention_block(block_class: Type[torch.nn.Module]) -> Type[torch.nn.Module]:
+
+    class TokenFlowBlock(block_class):
+
+        def forward(
+            self,
+            hidden_states,
+            attention_mask=None,
+            encoder_hidden_states=None,
+            encoder_attention_mask=None,
+            timestep=None,
+            cross_attention_kwargs=None,
+            class_labels=None,
+        ) -> torch.Tensor:
+            
+            batch_size, sequence_length, dim = hidden_states.shape
+            n_frames = batch_size // 3
+            mid_idx = n_frames // 2
+            hidden_states = hidden_states.view(3, n_frames, sequence_length, dim)
+
+            if self.use_ada_layer_norm:
+                norm_hidden_states = self.norm1(hidden_states, timestep)
+            elif self.use_ada_layer_norm_zero:
+                norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
+                    hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype
+                )
+            else:
+                norm_hidden_states = self.norm1(hidden_states)
+
+            norm_hidden_states = norm_hidden_states.view(3, n_frames, sequence_length, dim)
+            if self.pivotal_pass:
+                self.pivot_hidden_states = norm_hidden_states
+            else:
+                idx1 = []
+                idx2 = [] 
+                batch_idxs = [self.batch_idx]
+                if self.batch_idx > 0:
+                    batch_idxs.append(self.batch_idx - 1)
+                
+                sim = batch_cosine_sim(norm_hidden_states[0].reshape(-1, dim),
+                                        self.pivot_hidden_states[0][batch_idxs].reshape(-1, dim))
+                if len(batch_idxs) == 2:
+                    sim1, sim2 = sim.chunk(2, dim=1)
+                    # sim: n_frames * seq_len, len(batch_idxs) * seq_len
+                    idx1.append(sim1.argmax(dim=-1))  # n_frames * seq_len
+                    idx2.append(sim2.argmax(dim=-1))  # n_frames * seq_len
+                else:
+                    idx1.append(sim.argmax(dim=-1))
+                idx1 = torch.stack(idx1 * 3, dim=0) # 3, n_frames * seq_len
+                idx1 = idx1.squeeze(1)
+                if len(batch_idxs) == 2:
+                    idx2 = torch.stack(idx2 * 3, dim=0) # 3, n_frames * seq_len
+                    idx2 = idx2.squeeze(1)
+
+            # 1. Self-Attention
+            cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
+            if self.pivotal_pass:
+                # norm_hidden_states.shape = 3, n_frames * seq_len, dim
+                self.attn_output = self.attn1(
+                        norm_hidden_states.view(batch_size, sequence_length, dim),
+                        encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,
+                        **cross_attention_kwargs,
+                    )
+                # 3, n_frames * seq_len, dim - > 3 * n_frames, seq_len, dim
+                self.kf_attn_output = self.attn_output 
+            else:
+                batch_kf_size, _, _ = self.kf_attn_output.shape
+                self.attn_output = self.kf_attn_output.view(3, batch_kf_size // 3, sequence_length, dim)[:,
+                                   batch_idxs]  # 3, n_frames, seq_len, dim --> 3, len(batch_idxs), seq_len, dim
+            if self.use_ada_layer_norm_zero:
+                self.attn_output = gate_msa.unsqueeze(1) * self.attn_output
+
+            # gather values from attn_output, using idx as indices, and get a tensor of shape 3, n_frames, seq_len, dim
+            if not self.pivotal_pass:
+                if len(batch_idxs) == 2:
+                    attn_1, attn_2 = self.attn_output[:, 0], self.attn_output[:, 1]
+                    attn_output1 = attn_1.gather(dim=1, index=idx1.unsqueeze(-1).repeat(1, 1, dim))
+                    attn_output2 = attn_2.gather(dim=1, index=idx2.unsqueeze(-1).repeat(1, 1, dim))
+
+                    s = torch.arange(0, n_frames).to(idx1.device) + batch_idxs[0] * n_frames
+                    # distance from the pivot
+                    p1 = batch_idxs[0] * n_frames + n_frames // 2
+                    p2 = batch_idxs[1] * n_frames + n_frames // 2
+                    d1 = torch.abs(s - p1)
+                    d2 = torch.abs(s - p2)
+                    # weight
+                    w1 = d2 / (d1 + d2)
+                    w1 = torch.sigmoid(w1)
+                    
+                    w1 = w1.unsqueeze(0).unsqueeze(-1).unsqueeze(-1).repeat(3, 1, sequence_length, dim)
+                    attn_output1 = attn_output1.view(3, n_frames, sequence_length, dim)
+                    attn_output2 = attn_output2.view(3, n_frames, sequence_length, dim)
+                    attn_output = w1 * attn_output1 + (1 - w1) * attn_output2
+                else:
+                    attn_output = self.attn_output[:,0].gather(dim=1, index=idx1.unsqueeze(-1).repeat(1, 1, dim))
+
+                attn_output = attn_output.reshape(
+                        batch_size, sequence_length, dim)  # 3 * n_frames, seq_len, dim
+            else:
+                attn_output = self.attn_output
+            hidden_states = hidden_states.reshape(batch_size, sequence_length, dim)  # 3 * n_frames, seq_len, dim
+            hidden_states = attn_output + hidden_states
+
+            if self.attn2 is not None:
+                norm_hidden_states = (
+                    self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)
+                )
+
+                # 2. Cross-Attention
+                attn_output = self.attn2(
+                    norm_hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    attention_mask=encoder_attention_mask,
+                    **cross_attention_kwargs,
+                )
+                hidden_states = attn_output + hidden_states
+
+            # 3. Feed-forward
+            norm_hidden_states = self.norm3(hidden_states)
+
+            if self.use_ada_layer_norm_zero:
+                norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+
+
+            ff_output = self.ff(norm_hidden_states)
+
+            if self.use_ada_layer_norm_zero:
+                ff_output = gate_mlp.unsqueeze(1) * ff_output
+
+            hidden_states = ff_output + hidden_states
+
+            return hidden_states
+
+    return TokenFlowBlock
+
+
+def set_tokenflow(
+        model: torch.nn.Module):
+    """
+    Sets the tokenflow attention blocks in a model.
+    """
+
+    for _, module in model.named_modules():
+        if isinstance_str(module, "BasicTransformerBlock"):
+            make_tokenflow_block_fn = make_tokenflow_attention_block 
+            module.__class__ = make_tokenflow_block_fn(module.__class__)
+
+            # Something needed for older versions of diffusers
+            if not hasattr(module, "use_ada_layer_norm_zero"):
+                module.use_ada_layer_norm = False
+                module.use_ada_layer_norm_zero = False
+
+    return model

utils.py

@@ -0,0 +1,121 @@
+from pathlib import Path
+from PIL import Image
+import torch
+import yaml
+import math
+
+import torchvision.transforms as T
+from torchvision.io import read_video,write_video
+import os
+import random
+import numpy as np
+from torchvision.io import write_video
+# from kornia.filters import joint_bilateral_blur
+from kornia.geometry.transform import remap
+from kornia.utils.grid import create_meshgrid
+import cv2
+
+def save_video_frames(video_path, img_size=(512,512)):
+    video, _, _ = read_video(video_path, output_format="TCHW")
+    # rotate video -90 degree if video is .mov format. this is a weird bug in torchvision
+    if video_path.endswith('.mov'):
+        video = T.functional.rotate(video, -90)
+    video_name = Path(video_path).stem
+    os.makedirs(f'data/{video_name}', exist_ok=True)
+    for i in range(len(video)):
+        ind = str(i).zfill(5)
+        image = T.ToPILImage()(video[i])
+        image_resized = image.resize((img_size),  resample=Image.Resampling.LANCZOS)
+        image_resized.save(f'data/{video_name}/{ind}.png')
+        
+def video_to_frames(video_path, img_size=(512,512)):
+    video, _, _ = read_video(video_path, output_format="TCHW")
+    # rotate video -90 degree if video is .mov format. this is a weird bug in torchvision
+    if video_path.endswith('.mov'):
+        video = T.functional.rotate(video, -90)
+    video_name = Path(video_path).stem
+    # os.makedirs(f'data/{video_name}', exist_ok=True)
+    frames = []
+    for i in range(len(video)):
+        ind = str(i).zfill(5)
+        image = T.ToPILImage()(video[i])
+        image_resized = image.resize((img_size),  resample=Image.Resampling.LANCZOS)
+        # image_resized.save(f'data/{video_name}/{ind}.png')
+        frames.append(image_resized)
+    return frames
+
+def add_dict_to_yaml_file(file_path, key, value):
+    data = {}
+
+    # If the file already exists, load its contents into the data dictionary
+    if os.path.exists(file_path):
+        with open(file_path, 'r') as file:
+            data = yaml.safe_load(file)
+
+    # Add or update the key-value pair
+    data[key] = value
+
+    # Save the data back to the YAML file
+    with open(file_path, 'w') as file:
+        yaml.dump(data, file)
+        
+def isinstance_str(x: object, cls_name: str):
+    """
+    Checks whether x has any class *named* cls_name in its ancestry.
+    Doesn't require access to the class's implementation.
+    
+    Useful for patching!
+    """
+
+    for _cls in x.__class__.__mro__:
+        if _cls.__name__ == cls_name:
+            return True
+    
+    return False
+
+
+def batch_cosine_sim(x, y):
+    if type(x) is list:
+        x = torch.cat(x, dim=0)
+    if type(y) is list:
+        y = torch.cat(y, dim=0)
+    x = x / x.norm(dim=-1, keepdim=True)
+    y = y / y.norm(dim=-1, keepdim=True)
+    similarity = x @ y.T
+    return similarity
+
+
+def load_imgs(data_path, n_frames, device='cuda', pil=False):
+    imgs = []
+    pils = []
+    for i in range(n_frames):
+        img_path = os.path.join(data_path, "%05d.jpg" % i)
+        if not os.path.exists(img_path):
+            img_path = os.path.join(data_path, "%05d.png" % i)
+        img_pil = Image.open(img_path)
+        pils.append(img_pil)
+        img = T.ToTensor()(img_pil).unsqueeze(0)
+        imgs.append(img)
+    if pil:
+        return torch.cat(imgs).to(device), pils
+    return torch.cat(imgs).to(device)
+
+
+def save_video(raw_frames, save_path, fps=10):
+    video_codec = "libx264"
+    video_options = {
+        "crf": "18",  # Constant Rate Factor (lower value = higher quality, 18 is a good balance)
+        "preset": "slow",  # Encoding preset (e.g., ultrafast, superfast, veryfast, faster, fast, medium, slow, slower, veryslow)
+    }
+
+    frames = (raw_frames * 255).to(torch.uint8).cpu().permute(0, 2, 3, 1)
+    write_video(save_path, frames, fps=fps, video_codec=video_codec, options=video_options)
+
+
+def seed_everything(seed):
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+
+