diff --git a/animation/TokenFlow-WebUI/README.md b/animation/TokenFlow-WebUI/README.md deleted file mode 100644 index 8439116..0000000 --- a/animation/TokenFlow-WebUI/README.md +++ /dev/null @@ -1,33 +0,0 @@ -# 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 deleted file mode 100644 index 214c482..0000000 --- a/animation/TokenFlow-WebUI/app.py +++ /dev/null @@ -1,381 +0,0 @@ -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 deleted file mode 100644 index 9dbfb98..0000000 Binary files a/animation/TokenFlow-WebUI/examples/cutting_bread.mp4 and /dev/null differ diff --git a/animation/TokenFlow-WebUI/examples/rocket_kittens.mp4 b/animation/TokenFlow-WebUI/examples/rocket_kittens.mp4 deleted file mode 100644 index b980685..0000000 Binary files a/animation/TokenFlow-WebUI/examples/rocket_kittens.mp4 and /dev/null differ diff --git a/animation/TokenFlow-WebUI/examples/running_dog.mp4 b/animation/TokenFlow-WebUI/examples/running_dog.mp4 deleted file mode 100644 index 2fc9e88..0000000 Binary files a/animation/TokenFlow-WebUI/examples/running_dog.mp4 and /dev/null differ diff --git a/animation/TokenFlow-WebUI/examples/wolf.mp4 b/animation/TokenFlow-WebUI/examples/wolf.mp4 deleted file mode 100644 index 66ae4af..0000000 Binary files a/animation/TokenFlow-WebUI/examples/wolf.mp4 and /dev/null differ diff --git a/animation/TokenFlow-WebUI/examples/woman-running.mp4 b/animation/TokenFlow-WebUI/examples/woman-running.mp4 deleted file mode 100644 index 5db2cfb..0000000 Binary files a/animation/TokenFlow-WebUI/examples/woman-running.mp4 and /dev/null differ diff --git a/animation/TokenFlow-WebUI/preprocess_utils.py b/animation/TokenFlow-WebUI/preprocess_utils.py deleted file mode 100644 index b1ff9ec..0000000 --- a/animation/TokenFlow-WebUI/preprocess_utils.py +++ /dev/null @@ -1,333 +0,0 @@ -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 deleted file mode 100644 index b900454..0000000 --- a/animation/TokenFlow-WebUI/requirements.txt +++ /dev/null @@ -1,20 +0,0 @@ ---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 deleted file mode 100644 index f934522..0000000 --- a/animation/TokenFlow-WebUI/style.css +++ /dev/null @@ -1,91 +0,0 @@ -/* -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 deleted file mode 100644 index 74b940b..0000000 --- a/animation/TokenFlow-WebUI/tokenflow_pnp.py +++ /dev/null @@ -1,364 +0,0 @@ -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 deleted file mode 100644 index 6b76ec6..0000000 --- a/animation/TokenFlow-WebUI/tokenflow_utils.py +++ /dev/null @@ -1,448 +0,0 @@ -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 deleted file mode 100644 index 63bee1c..0000000 --- a/animation/TokenFlow-WebUI/utils.py +++ /dev/null @@ -1,121 +0,0 @@ -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) - -