vidtome/patch.py

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