diff --git a/animation/TokenFlow-WebUI/README.md b/animation/TokenFlow-WebUI/README.md new file mode 100644 index 0000000..8439116 --- /dev/null +++ b/animation/TokenFlow-WebUI/README.md @@ -0,0 +1,33 @@ +# TokenFlow-WebUI + +## Installation + +### Using Anaconda on Windows/Linux (Recommended) + +1. **Clone the repository:** + ```bash + git clone https://github.com/justinjohn0306/TokenFlow-WebUI.git + cd TokenFlow-WebUI + ``` + +2. **Create a new Anaconda environment named `tokenflow` with Python 3.9:** + ```bash + conda create -n tokenflow python=3.9 + ``` + +3. **Activate the `tokenflow` environment:** + ```bash + conda activate tokenflow + ``` + +4. **Install the required packages:** + ```bash + pip install -r requirements.txt + ``` + +## Running the App + +After you've installed all the required packages, you can launch the Gradio app by running the following command: + +```bash +python app.py diff --git a/animation/TokenFlow-WebUI/app.py b/animation/TokenFlow-WebUI/app.py new file mode 100644 index 0000000..214c482 --- /dev/null +++ b/animation/TokenFlow-WebUI/app.py @@ -0,0 +1,381 @@ +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") +if torch.cuda.is_available(): + device = "cuda" +elif torch.backends.mps.is_available(): + device = "mps" +else: + device = "cpu" +model_id = "stabilityai/stable-diffusion-2-1-base" + +to = torch.float16 if device == 'cuda' else torch.float32 + +# components for the Preprocessor +scheduler = DDIMScheduler.from_pretrained(model_id, subfolder="scheduler") +vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae", revision="fp16", + torch_dtype=to).to(device) +tokenizer = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer") +text_encoder = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder", revision="fp16", + torch_dtype=to).to(device) +unet = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet", revision="fp16", + torch_dtype=to).to(device) + +# pipe for TokenFlow +tokenflow_pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=to).to(device) +if device == "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 = """ +
+

+ TokenFlow - Temporally consistent video editing +

+[Project page], [GitHub], [Paper] +
+""" + + + +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=200, 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=25, 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() diff --git a/animation/TokenFlow-WebUI/examples/cutting_bread.mp4 b/animation/TokenFlow-WebUI/examples/cutting_bread.mp4 new file mode 100644 index 0000000..9dbfb98 Binary files /dev/null and b/animation/TokenFlow-WebUI/examples/cutting_bread.mp4 differ diff --git a/animation/TokenFlow-WebUI/examples/rocket_kittens.mp4 b/animation/TokenFlow-WebUI/examples/rocket_kittens.mp4 new file mode 100644 index 0000000..b980685 Binary files /dev/null and b/animation/TokenFlow-WebUI/examples/rocket_kittens.mp4 differ diff --git a/animation/TokenFlow-WebUI/examples/running_dog.mp4 b/animation/TokenFlow-WebUI/examples/running_dog.mp4 new file mode 100644 index 0000000..2fc9e88 Binary files /dev/null and b/animation/TokenFlow-WebUI/examples/running_dog.mp4 differ diff --git a/animation/TokenFlow-WebUI/examples/wolf.mp4 b/animation/TokenFlow-WebUI/examples/wolf.mp4 new file mode 100644 index 0000000..66ae4af Binary files /dev/null and b/animation/TokenFlow-WebUI/examples/wolf.mp4 differ diff --git a/animation/TokenFlow-WebUI/examples/woman-running.mp4 b/animation/TokenFlow-WebUI/examples/woman-running.mp4 new file mode 100644 index 0000000..5db2cfb Binary files /dev/null and b/animation/TokenFlow-WebUI/examples/woman-running.mp4 differ diff --git a/animation/TokenFlow-WebUI/preprocess_utils.py b/animation/TokenFlow-WebUI/preprocess_utils.py new file mode 100644 index 0000000..b1ff9ec --- /dev/null +++ b/animation/TokenFlow-WebUI/preprocess_utils.py @@ -0,0 +1,333 @@ +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.to = torch.float16 if self.device == 'cuda' else torch.float32 + 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=self.to).to(self.device) + control_pipe = StableDiffusionControlNetPipeline.from_pretrained( + "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=self.to + ).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 + + if self.device == "cuda": + 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(self.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(self.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(self.to) + + @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(self.to) + 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(self.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(self.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(self.to).to(self.device) + # encode to latents + latents = self.encode_imgs(frames, deterministic=True).to(self.to).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 + + diff --git a/animation/TokenFlow-WebUI/requirements.txt b/animation/TokenFlow-WebUI/requirements.txt new file mode 100644 index 0000000..b900454 --- /dev/null +++ b/animation/TokenFlow-WebUI/requirements.txt @@ -0,0 +1,20 @@ +--extra-index-url https://download.pytorch.org/whl/cu118 + +pillow +diffusers==0.20.0 +ftfy +transformers +opencv-python +tqdm +numpy +pyyaml +accelerate +xformers +tensorboard +kornia +av +gradio==3.41.2 +torch==2.0.1+cu118; sys_platform != 'darwin' +torch==2.0.1; sys_platform == 'darwin' +torchvision==0.15.2+cu118; sys_platform != 'darwin' +torchvision==0.15.2; sys_platform == 'darwin' diff --git a/animation/TokenFlow-WebUI/style.css b/animation/TokenFlow-WebUI/style.css new file mode 100644 index 0000000..f934522 --- /dev/null +++ b/animation/TokenFlow-WebUI/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; +} \ No newline at end of file diff --git a/animation/TokenFlow-WebUI/tokenflow_pnp.py b/animation/TokenFlow-WebUI/tokenflow_pnp.py new file mode 100644 index 0000000..74b940b --- /dev/null +++ b/animation/TokenFlow-WebUI/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"] + self.to = torch.float16 if self.device == 'cuda' else torch.float32 + + 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(self.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(self.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(to).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(self.to).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(self.to).to(self.device) + # get noise + eps = self.get_ddim_eps(latents, range(self.config["n_frames"])).to(self.to).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) diff --git a/animation/TokenFlow-WebUI/tokenflow_utils.py b/animation/TokenFlow-WebUI/tokenflow_utils.py new file mode 100644 index 0000000..6b76ec6 --- /dev/null +++ b/animation/TokenFlow-WebUI/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 diff --git a/animation/TokenFlow-WebUI/utils.py b/animation/TokenFlow-WebUI/utils.py new file mode 100644 index 0000000..63bee1c --- /dev/null +++ b/animation/TokenFlow-WebUI/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) + +