import os import math import time from typing import Type, Dict, Any, Tuple, Callable import numpy as np from einops import rearrange, repeat import torch import torch.nn.functional as F from . import merge from .utils import isinstance_str, init_generator, join_frame, split_frame, func_warper, join_warper, split_warper def compute_merge(module: torch.nn.Module, x: torch.Tensor, tome_info: Dict[str, Any]) -> Tuple[Callable, ...]: H, original_w = tome_info["size"] # original_tokens = original_h * original_w # downsample = int(math.ceil(math.sqrt(original_tokens // x.shape[1]))) downsample = tome_info["args"]["downsample"] args = tome_info["args"] # generator = module.generator # Frame Number and Token Number fsize = x.shape[0] // args["batch_size"] tsize = x.shape[1] # Merge tokens in high resolution layers # print(f"[INFO] {args['current_step']} downsample {downsample} time") mid = x.shape[0] // 2 ''' visualize token correspondence ''' label = args["label"].split('_') # os.makedirs(os.path.join("token_0204_dis", str(args["current_step"])), exist_ok=True) # os.makedirs(os.path.join("token_0204_dis", str(args["current_step"]), label[0]), exist_ok=True) # out = os.path.join("token_0204_dis", str(args["current_step"]), label[0], f"correspondence_{label[1]}_{downsample}.png") # merge.visualize_correspondence(x[0][None], x[mid][None], ratio=0.2, H=H, out=out) # corres_dir = "corres_no_dis_4" # os.makedirs(corres_dir, exist_ok=True) # if downsample == 1 and label[0] == "unet" and label[1] == "down": # # merge.visualize_flow_correspondence(x[3][None], x[mid][None], flow=args["controller"].step_store["flows"][3], flow_confid=args["controller"].step_store["flow_confids"][3], \ # # ratio=0.1, H=(64//downsample), out=f"flow_{label[1]}_{args['current_step']}.png") # files = os.listdir(corres_dir) # files = [f for f in files if f.startswith(f"{args['current_step']}")] # print(files) # cnt = len(files) # # if files: # # cnt = int(files[-1].split('_')[1].split('.')[0]) + 1 # # else: # # cnt = 0 # path = os.path.join(corres_dir, f"{args['current_step']}_{cnt}.png") # # merge.visualize_cosine_correspondence(x[3][None], x[mid][None], flow=args["controller"].step_store["flows"][3], flow_confid=args["controller"].step_store["flow_confids"][3], \ # # ratio=0.1, H=(64//downsample), out=path, controller=args["controller"]) # merge.visualize_correspondence(x[1][None], x[mid][None], flow=args["controller"].step_store["flows"][1], flow_confid=args["controller"].step_store["flow_confids"][1], \ # ratio=0.1, H=(64//downsample), out=path, controller=args["controller"]) # #ratio=args["local_merge_ratio"], H=(64//downsample), out=f"flow_correspondence.png") ''' visulaize token correspondence ended ''' if downsample <= args["max_downsample"] and downsample > args["min_downsample"]: # print(f"[INFO] downsample: {args['min_downsample']} < {downsample} <= {args['max_downsample']} token shape: {x.shape} H: {H}") if args["generator"] is None: args["generator"] = init_generator(x.device) # module.generator = module.generator.manual_seed(123) elif args["generator"].device != x.device: args["generator"] = init_generator(x.device, fallback=args["generator"]) # Local Token Merging! local_tokens = join_frame(x, fsize) m_ls = [join_warper(fsize)] u_ls = [split_warper(fsize)] unm = 0 curF = fsize # Recursive merge multi-frame tokens into one set. Such as 4->1 for 4 frames and 8->2->1 for 8 frames when target stride is 4. while curF > 1: current_step = args["current_step"] if args["controller"] is not None: controller, total_step = args["controller"], args["controller"].total_step else: controller, total_step = None, 1000 if controller is not None and label[0] == "unet" and label[1] == "down": print(f"[INFO] flow merge @ {label[0]} {label[1]} {downsample}") start = time.time() m, u, ret_dict = merge.bipartite_soft_matching_randframe( local_tokens, curF, args["local_merge_ratio"], unm, generator=args["generator"], target_stride=x.shape[0], align_batch=args["align_batch"], H=H, flow_merge=True, controller=controller, ) else: m, u, ret_dict = merge.bipartite_soft_matching_randframe( local_tokens, curF, args["local_merge_ratio"], unm, generator=args["generator"], target_stride=x.shape[0], align_batch=args["align_batch"], H=H, flow_merge=False, controller=controller, ) # if controller is not None and label[1] == "up" and \ # controller.merge_period[0] > 0 and \ # (current_step + 5) >= min(controller.ToMe_period[1], controller.merge_period[0]) * total_step: # # or current_step == int(controller.ToMe_period[1] * total_step)): # print(f"[INFO] setting controller merge @ step {current_step} {label} {downsample}") # # ret_dict["scores"].repeat(1, 4, 4) # # import time # # start = time.time() # scores = ret_dict["scores"] # if downsample > 1: # scores = rearrange(scores, "1 (b h1 w1) (h2 w2) -> b h1 w1 h2 w2", h1=H, h2=H, b=fsize-1) # scores = scores.repeat_interleave(downsample, dim=-1).repeat_interleave(downsample, dim=-2) # scores = scores.repeat_interleave(downsample, dim=1).repeat_interleave(downsample, dim=2) # scores = rearrange(scores, "b h1 w1 h2 w2 -> 1 (b h1 w1) (h2 w2)") # # print(f"[INFO] repeat time {time.time() - start}") # # import ipdb; ipdb.set_trace() # # merge.visualize_correspondence_score(x[0][None], x[mid][None], score=ret_dict["scores"][:,:tsize], ratio=0.5, H=H, out="latent_correspondence_1.png") # # merge.visualize_correspondence_score(x[0][None], x[mid][None], score=controller.step_store["corres_scores"][:,:tsize], ratio=0.5, H=H, out="latent_correspondence_1.png") # if controller.step_store["corres_scores"] is None: # controller.step_store["corres_scores"] = scores # else: # controller.step_store["corres_scores"] += scores unm += ret_dict["unm_num"] m_ls.append(m) u_ls.append(u) local_tokens = m(local_tokens) # assert (x.shape[1] - unm) % tsize == 0 # Total token number = current frame number * per-frame token number + unmerged token number curF = (local_tokens.shape[1] - unm) // tsize # print(f"[INFO] curF {curF}") merged_tokens = local_tokens # Global Token Merging! if args["merge_global"]: if hasattr(module, "global_tokens") and module.global_tokens is not None: # Merge local tokens with global tokens. Randomly determine merging destination. if torch.rand(1, generator=args["generator"], device=args["generator"].device) > args["global_rand"]: src_len = local_tokens.shape[1] tokens = torch.cat( [local_tokens, module.global_tokens.to(local_tokens)], dim=1) local_chunk = 0 else: src_len = module.global_tokens.shape[1] tokens = torch.cat( [module.global_tokens.to(local_tokens), local_tokens], dim=1) local_chunk = 1 m, u, _ = merge.bipartite_soft_matching_2s( tokens, src_len, args["global_merge_ratio"], args["align_batch"], unmerge_chunk=local_chunk) merged_tokens = m(tokens) # print(f"[INFO] global merging {local_tokens.shape} {tokens.shape} {merged_tokens.shape}") # import ipdb; ipdb.set_trace() m_ls.append(m) u_ls.append(u) # Update global tokens with unmerged local tokens. There should be a better way to do this. module.global_tokens = u(merged_tokens).detach().clone().cpu() else: module.global_tokens = local_tokens.detach().clone().cpu() m = func_warper(m_ls) u = func_warper(u_ls[::-1]) else: m, u = (merge.do_nothing, merge.do_nothing) merged_tokens = x # if args["current_step"] >= 30: # x_ = u(m(x)) # print(f"[INFO] {label[0]} {label[1]} {downsample}") # for i, j in zip(x, x_): # print(f"[INFO] mean {torch.mean(i).item()} {torch.mean(j).item()}") # print(f"[INFO] std {torch.std(i).item()} {torch.std(j).item()}") # import ipdb; ipdb.set_trace() # Return merge op, unmerge op, and merged tokens. return m, u, merged_tokens def PCA_token(token: torch.Tensor, token_h=64, n=3): from sklearn.decomposition import PCA import cv2 pca = PCA(n_components=n) # reduce to 2 dimensions # Fit the PCA model to your data and apply the dimensionality reduction token = pca.fit_transform(token[0].cpu()) # import ipdb; ipdb.set_trace() token = rearrange(token, '(h w) c -> h w c', h=token_h) token = (token - token.min()) / (token.max() - token.min()) token = (np.clip(token, 0, 1) * 255).astype(np.uint8) cv2.imwrite(f'token.png', token) return token from utils.flow_utils import flow_warp def warp_token(module: torch.nn.Module, x: torch.Tensor, tome_info: Dict[str, Any]) -> Tuple[Callable, ...]: original_h, original_w = tome_info["size"] original_tokens = original_h * original_w downsample = int(math.ceil(math.sqrt(original_tokens // x.shape[1]))) args = tome_info["args"] # generator = module.generator # Frame Number and Token Number fsize = x.shape[0] // args["batch_size"] tsize = x.shape[1] # print(f"[INFO] token size {x.shape[1]}, latent size 64 x 120, downsample {downsample} time") # Merge tokens in high resolution layers total_step = 50 warp_period = (0, 1) if downsample <= args["max_downsample"] and x.shape[1] == 64 * 120: if args["current_step"] >= total_step * warp_period[0] and \ args["current_step"] <= total_step * warp_period[1]: mid = x.shape[0] // 2 x = rearrange(x, 'b (h w) c -> b c h w', h=64) # import ipdb; ipdb.set_trace() # mid_x = repeat(x[mid][None], 'b c h w -> (repeat b) c h w', repeat=x.shape[0]) for i in range(x.shape[0]): if i == mid: continue x[i] = flow_warp(x[mid][None], args["flows"][i], mode='nearest')[0] * args["occlusion_masks"][i] + \ (1 - args["occlusion_masks"][i]) * x[i] x = rearrange(x, 'b c h w -> b (h w) c', h=64) return x def make_tome_block(block_class: Type[torch.nn.Module]) -> Type[torch.nn.Module]: """ Make a patched class on the fly so we don't have to import any specific modules. This patch applies ToMe to the forward function of the block. """ class ToMeBlock(block_class): # Save for unpatching later _parent = block_class def _forward(self, x: torch.Tensor, context: torch.Tensor = None, label: str = None) -> torch.Tensor: # m_a, m_c, m_m, u_a, u_c, u_m = compute_merge( # self, x, self._tome_info) # print(f"[INFO] ~~~ ToMeblock ~~~ {label} ~~~") B, A, C = x.shape original_h, original_w = self._tome_info["size"] original_tokens = original_h * original_w downsample = int(math.ceil(math.sqrt(original_tokens // A))) # print(f"[INFO] downsample {downsample} time A {A} original_h {original_h} original_w {original_w}") self._tome_info["args"]["downsample"] = downsample H, W = original_h // downsample, original_w // downsample if self._tome_info["args"]["controller"] is None: non_pad_ratio_h, non_pad_ratio_w = 1, 1 print(f"[INFO] no padding removal") else: non_pad_ratio_h, non_pad_ratio_w = self._tome_info["args"]["controller"].non_pad_ratio padding_size_w = W - int(W * non_pad_ratio_w) padding_size_h = H - int(H * non_pad_ratio_h) padding_mask = torch.zeros((H, W), device=x.device, dtype=torch.bool) if padding_size_w: padding_mask[:, -padding_size_w:] = 1 if padding_size_h: padding_mask[-padding_size_h:, :] = 1 padding_mask = rearrange(padding_mask, 'h w -> (h w)') idx_buffer = torch.arange(A, device=x.device, dtype=torch.int64) non_pad_idx = idx_buffer[None, ~padding_mask, None] # pad_idx = idx_buffer[None, padding_mask, None] del idx_buffer, padding_mask x_non_pad = torch.gather(x, dim=1, index=non_pad_idx.expand(B, -1, C)) self._tome_info["args"]["label"] = label self._tome_info["size"] = (int(H * non_pad_ratio_h), int(W * non_pad_ratio_w)) # self._tome_info["non_pad_size"] = (int(H * non_pad_ratio_h), int(W * non_pad_ratio_w)) # print(f"[INFO] original shape {x.shape} removed padding shape {x_non_pad.shape}") m_a, u_a, merged_tokens = compute_merge( self, self.norm1(x_non_pad), self._tome_info) # current_step, total_step = self._tome_info["args"]["current_step"], self._tome_info["args"]["controller"].total_step # print(f'[INFO] {int(self._tome_info["args"]["controller"].ToMe_period[1] * total_step)} {current_step} {total_step}') # if downsample == 1 and label == "unet_up" and \ # self._tome_info["args"]["controller"].merge_period[0] > 0 and \ # (current_step >= self._tome_info["args"]["controller"].merge_period[0] * total_step \ # or current_step == int(self._tome_info["args"]["controller"].ToMe_period[1] * total_step)): # print(f"[INFO] setting controller merge @ step {self._tome_info['args']['current_step']}") # self._tome_info["args"]["controller"].set_merge(m_a, u_a) # m_a, u_a, merged_tokens = compute_merge( # self, self.norm1(x), self._tome_info) # This is where the meat of the computation happens # test = u_a(self.attn1(m_a(self.norm1(x)), context=context if self.disable_self_attn else None)) # import ipdb; ipdb.set_trace() # x = u_a(merged_tokens) ''' global merging ''' if self._tome_info["args"]["controller"] is None: print(f"[INFO] local + global merging ... ") x_non_pad = u_a(self.attn1(merged_tokens, context=context if self.disable_self_attn else None)) + x_non_pad else: x_non_pad = u_a(self.attn1(m_a(self.norm1(x_non_pad)), context=context if self.disable_self_attn else None)) + x_non_pad # print(label, downsample, self.disable_self_attn) # x = u_a(self.attn1(m_a(self.norm1(x)), # context=context if self.disable_self_attn else None)) + x # attn_output = self.attn1(self.norm1(x), context=context if self.disable_self_attn else None) # attn_output = warp_token(self, attn_output, self._tome_info) # x = attn_output + x # attn_out = self.attn2(self.norm2(x), context=context) # m_a, u_a, merged_tokens = compute_merge( # self, attn_out, self._tome_info) # x = u_a(m_a(attn_out)) + x # attn_output = self.attn2(self.norm2(x), context=context) # attn_output = warp_token(self, attn_output, self._tome_info) # x = attn_output + x x_non_pad = self.attn2(self.norm2(x_non_pad), context=context) + x_non_pad x_non_pad = self.ff(self.norm3(x_non_pad)) + x_non_pad x.scatter_(dim=1, index=non_pad_idx.expand(B, -1, C), src=x_non_pad) del x_non_pad self._tome_info["size"] = (original_h, original_w) torch.cuda.empty_cache() # x = self.attn2(self.norm2(x), context=context) + x # x = self.ff(self.norm3(x)) + x return x return ToMeBlock def make_diffusers_tome_block(block_class: Type[torch.nn.Module]) -> Type[torch.nn.Module]: """ Make a patched class for a diffusers model. This patch applies ToMe to the forward function of the block. """ class ToMeBlock(block_class): # Save for unpatching later _parent = 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: 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) # Merge input tokens m_a, u_a, merged_tokens = compute_merge( self, norm_hidden_states, self._tome_info) norm_hidden_states = merged_tokens # 1. Self-Attention cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} # tt = time.time() attn_output = self.attn1( norm_hidden_states, encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None, attention_mask=attention_mask, **cross_attention_kwargs, ) # print(time.time() - tt) if self.use_ada_layer_norm_zero: attn_output = gate_msa.unsqueeze(1) * attn_output # Unmerge output tokens attn_output = u_a(attn_output) 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 ToMeBlock def hook_tome_model(model: torch.nn.Module): """ Adds a forward pre hook to get the image size. This hook can be removed with remove_patch. """ def hook(module, args): # print(args[0].shape) module._tome_info["size"] = (args[0].shape[2], args[0].shape[3]) return None model._tome_info["hooks"].append(model.register_forward_pre_hook(hook)) def hook_tome_module(module: torch.nn.Module): """ Adds a forward pre hook to initialize random number generator. All modules share the same generator state to keep their randomness in VidToMe consistent in one pass. This hook can be removed with remove_patch. """ def hook(module, args): # import ipdb; ipdb.set_trace() if not hasattr(module, "generator"): module.generator = init_generator(args[0].device) elif module.generator.device != args[0].device: module.generator = init_generator( args[0].device, fallback=module.generator) else: return None # module.generator = module.generator.manual_seed(module._tome_info["args"]["seed"]) return None module._tome_info["hooks"].append(module.register_forward_pre_hook(hook)) def apply_patch( model: torch.nn.Module, local_merge_ratio: float = 0.9, merge_global: bool = False, global_merge_ratio = 0.8, max_downsample: int = 2, min_downsample: int = 0, seed: int = 123, batch_size: int = 2, include_control: bool = False, align_batch: bool = False, target_stride: int = 4, global_rand=0.5): """ Patches a stable diffusion model with VidToMe. Apply this to the highest level stable diffusion object (i.e., it should have a .model.diffusion_model). Important Args: - model: A top level Stable Diffusion module to patch in place. Should have a ".model.diffusion_model" - local_merge_ratio: The ratio of tokens to merge locally. I.e., 0.9 would merge 90% src tokens. If there are 4 frames in a chunk (3 src, 1 dst), the compression ratio will be 1.3 / 4.0. And the largest compression ratio is 0.25 (when local_merge_ratio = 1.0). Higher values result in more consistency, but with more visual quality loss. - merge_global: Whether or not to include global token merging. - global_merge_ratio: The ratio of tokens to merge locally. I.e., 0.8 would merge 80% src tokens. When find significant degradation in video quality. Try to lower the value. Args to tinker with if you want: - max_downsample [1, 2, 4, or 8]: Apply VidToMe to layers with at most this amount of downsampling. E.g., 1 only applies to layers with no downsampling (4/15) while 8 applies to all layers (15/15). I recommend a value of 1 or 2. - seed: Manual random seed. - batch_size: Video batch size. Number of video chunks in one pass. When processing one video, it should be 2 (cond + uncond) or 3 (when using PnP, source + cond + uncond). - include_control: Whether or not to patch ControlNet model. - align_batch: Whether or not to align similarity matching maps of samples in the batch. It should be True when using PnP as control. - target_stride: Stride between target frames. I.e., when target_stride = 4, there is 1 target frame in any 4 consecutive frames. - global_rand: Probability in global token merging src/dst split. Global tokens are always src when global_rand = 1.0 and always dst when global_rand = 0.0 . """ # Make sure the module is not currently patched remove_patch(model) is_diffusers = isinstance_str( model, "DiffusionPipeline") or isinstance_str(model, "ModelMixin") if not is_diffusers: if (not hasattr(model, "model") or not hasattr(model.model, "diffusion_model")) \ and not hasattr(model, "unet"): # Provided model not supported raise RuntimeError( "Provided model was not a Stable Diffusion / Latent Diffusion model, as expected.") else: diffusion_model = model.unet if hasattr(model, "unet") else model.model.diffusion_model else: # Supports "pipe.unet" and "unet" diffusion_model = model.unet if hasattr(model, "unet") else model if isinstance_str(model, "StableDiffusionControlNetPipeline") and include_control: diffusion_models = [diffusion_model, model.controlnet] else: diffusion_models = [diffusion_model] if not is_diffusers and hasattr(model, "controlnet"): diffusion_models = [diffusion_model, model.controlnet] for diffusion_model in diffusion_models: diffusion_model._tome_info = { "size": None, "hooks": [], "args": { "max_downsample": max_downsample, "min_downsample": min_downsample, "generator": None, "seed": seed, "batch_size": batch_size, "align_batch": align_batch, "merge_global": merge_global, "global_merge_ratio": global_merge_ratio, "local_merge_ratio": local_merge_ratio, "global_rand": global_rand, "target_stride": target_stride, "current_step": 0, "frame_ids": [0], "flows": None, "occlusion_masks": None, "flow_confids": None, "label": "", "downsample": 1, "non_pad_size": (0, 0), "controller": None, } } hook_tome_model(diffusion_model) for name, 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") and "down_blocks" not in name: # print(module.__class__.__name__) if isinstance_str(module, "BasicTransformerBlock"): make_tome_block_fn = make_diffusers_tome_block if is_diffusers else make_tome_block module.__class__ = make_tome_block_fn(module.__class__) module._tome_info = diffusion_model._tome_info hook_tome_module(module) # Something introduced in SD 2.0 (LDM only) if not hasattr(module, "disable_self_attn") and not is_diffusers: module.disable_self_attn = False # Something needed for older versions of diffusers if not hasattr(module, "use_ada_layer_norm_zero") and is_diffusers: module.use_ada_layer_norm = False module.use_ada_layer_norm_zero = False # import ipdb; ipdb.set_trace() return model def remove_patch(model: torch.nn.Module): """ Removes a patch from a ToMe Diffusion module if it was already patched. """ # For diffusers modelu = model.unet if hasattr(model, "unet") else model model_ls = [modelu] if hasattr(model, "controlnet"): model_ls.append(model.controlnet) for model in model_ls: for _, module in model.named_modules(): if hasattr(module, "_tome_info"): for hook in module._tome_info["hooks"]: hook.remove() module._tome_info["hooks"].clear() if module.__class__.__name__ == "ToMeBlock": module.__class__ = module._parent return model def update_patch(model: torch.nn.Module, **kwargs): """ Update arguments in patched modules """ # For diffusers model0 = model.unet if hasattr(model, "unet") else model model_ls = [model0] if hasattr(model, "controlnet"): model_ls.append(model.controlnet) for model in model_ls: for _, module in model.named_modules(): if hasattr(module, "_tome_info"): for k, v in kwargs.items(): # setattr(module, k, v) if k in module._tome_info["args"]: module._tome_info["args"][k] = v # print(f"[INFO] update {k} to {v} in {module.__class__.__name__}") return model def collect_from_patch(model: torch.nn.Module, attr="tome"): """ Collect attributes in patched modules """ # For diffusers model0 = model.unet if hasattr(model, "unet") else model model_ls = [model0] if hasattr(model, "controlnet"): model_ls.append(model.controlnet) ret_dict = dict() for model in model_ls: for name, module in model.named_modules(): if hasattr(module, attr): res = getattr(module, attr) ret_dict[name] = res return ret_dict