update hf hub

app.py

@@ -2,35 +2,15 @@ import gradio as gr
 from dataclasses import dataclass
 import spaces
 import torch
-from tqdm import tqdm
+from huggingface_hub import hf_hub_download
 
-from src.utils import (
-    create_pipeline,
-    calculate_mask_sparsity,
-    ffn_linear_layer_pruning,
-    linear_layer_pruning,
-)
-from diffusers import StableDiffusionXLPipeline
+from diffusers import StableDiffusionXLPipeline, FluxPipeline
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
-def get_model_param_summary(model, verbose=False):
-    params_dict = dict()
-    overall_params = 0
-    for name, params in model.named_parameters():
-        num_params = params.numel()
-        overall_params += num_params
-        if verbose:
-            print(f"GPU Memory Requirement for {name}: {params} MiB")
-        params_dict.update({name: num_params})
-    params_dict.update({"overall": overall_params})
-    return params_dict
-
-
 @dataclass
 class GradioArgs:
-    ckpt: str = "./mask/ff.pt"
     seed: list = None
     prompt: str = None
     mix_precision: str = "bf16"
@@ -48,96 +28,44 @@ class GradioArgs:
     def __post_init__(self):
         if self.seed is None:
             self.seed = [44]
-        if self.ratio is None:
-            self.ratio = [0.68, 0.88]
-
-
-def prune_model(pipe, hookers):
-    # remove parameters in attention blocks
-    cross_attn_hooker = hookers[0]
-    for name in tqdm(cross_attn_hooker.hook_dict.keys(), desc="Pruning attention layers"):
-        if getattr(pipe, "unet", None):
-            module = pipe.unet.get_submodule(name)
-        else:
-            module = pipe.transformer.get_submodule(name)
-        lamb = cross_attn_hooker.lambs[cross_attn_hooker.lambs_module_names.index(name)]
-        assert module.heads == lamb.shape[0]
-        module = linear_layer_pruning(module, lamb)
-
-        parent_module_name, child_name = name.rsplit(".", 1)
-        if getattr(pipe, "unet", None):
-            parent_module = pipe.unet.get_submodule(parent_module_name)
-        else:
-            parent_module = pipe.transformer.get_submodule(parent_module_name)
-        setattr(parent_module, child_name, module)
-
-    # remove parameters in ffn blocks
-    ffn_hook = hookers[1]
-    for name in tqdm(ffn_hook.hook_dict.keys(), desc="Pruning on FFN linear lazer"):
-        if getattr(pipe, "unet", None):
-            module = pipe.unet.get_submodule(name)
-        else:
-            module = pipe.transformer.get_submodule(name)
-        lamb = ffn_hook.lambs[ffn_hook.lambs_module_names.index(name)]
-        module = ffn_linear_layer_pruning(module, lamb)
-
-        parent_module_name, child_name = name.rsplit(".", 1)
-        if getattr(pipe, "unet", None):
-            parent_module = pipe.unet.get_submodule(parent_module_name)
-        else:
-            parent_module = pipe.transformer.get_submodule(parent_module_name)
-        setattr(parent_module, child_name, module)
-
-    cross_attn_hooker.clear_hooks()
-    ffn_hook.clear_hooks()
-    return pipe
-
-
-def binary_mask_eval(args):
-    # load sdxl model
-    pipe = StableDiffusionXLPipeline.from_pretrained(
-        "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16
-    ).to("cpu")
-
-    torch_dtype = torch.bfloat16 if args.mix_precision == "bf16" else torch.float32
-    mask_pipe, hookers = create_pipeline(
-        pipe,
-        args.model,
-        "cpu",
-        torch_dtype,
-        args.ckpt,
-        binary=args.binary,
-        lambda_threshold=args.lambda_threshold,
-        epsilon=args.epsilon,
-        masking=args.masking,
-        return_hooker=True,
-        scope=args.scope,
-        ratio=args.ratio,
-    )
 
-    # # Print mask sparsity info
-    # threshold = None if args.binary else args.lambda_threshold
-    # threshold = None if args.scope is not None else threshold
-    # name = ["ff", "attn"]
-    # for n, hooker in zip(name, hookers):
-    #     total_num_heads, num_activate_heads, mask_sparsity = calculate_mask_sparsity(hooker, threshold)
-    #     print(f"model: {args.model}, {n} masking: {args.masking}")
-    #     print(
-    #         f"total num heads: {total_num_heads},"
-    #         + f"num activate heads: {num_activate_heads}, mask sparsity: {mask_sparsity}"
-    #     )
-
-    # Prune the model
-    pruned_pipe = prune_model(mask_pipe, hookers)
+
+def binary_mask_eval(args, model):
+    model = model.lower()
+    # load sdxl model
+    if model == "sdxl":
+        pruned_pipe = StableDiffusionXLPipeline.from_pretrained(
+            "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16
+        ).to("cpu")
+        pruned_pipe.unet = torch.load(
+            hf_hub_download(
+                "zhangyang-0123/EcoDiffPrunedModels", "model/sdxl/sdxl.pkl"
+            ),
+            map_location="cpu",
+        )
+    elif model == "flux":
+        pruned_pipe = FluxPipeline.from_pretrained(
+            "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16
+        ).to("cpu")
+        pruned_pipe.transformer = torch.load(
+            hf_hub_download(
+                "zhangyang-0123/EcoDiffPrunedModels", "model/flux/flux.pkl"
+            ),
+            map_location="cpu",
+        )
+        torch.cuda.empty_cache()
 
     # reload the original model
-    pipe = StableDiffusionXLPipeline.from_pretrained(
-        "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16
-    ).to("cpu")
+    if model == "sdxl":
+        pipe = StableDiffusionXLPipeline.from_pretrained(
+            "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16
+        ).to("cpu")
+    elif model == "flux":
+        pipe = FluxPipeline.from_pretrained(
+            "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16
+        ).to("cpu")
+        torch.cuda.empty_cache()
 
-    # get model param summary
-    print(f"original model param: {get_model_param_summary(pipe.unet)['overall']}")
-    print(f"pruned model param: {get_model_param_summary(pruned_pipe.unet)['overall']}")
     print("prune complete")
     return pipe, pruned_pipe
 
@@ -148,20 +76,27 @@ def generate_images(prompt, seed, steps, pipe, pruned_pipe):
     pruned_pipe.to("cuda")
     # Run the model and return images directly
     g_cpu = torch.Generator("cuda").manual_seed(seed)
-    original_image = pipe(prompt=prompt, generator=g_cpu, num_inference_steps=steps).images[0]
+    original_image = pipe(
+        prompt=prompt, generator=g_cpu, num_inference_steps=steps
+    ).images[0]
+    torch.cuda.empty_cache()
     g_cpu = torch.Generator("cuda").manual_seed(seed)
-    ecodiff_image = pruned_pipe(prompt=prompt, generator=g_cpu, num_inference_steps=steps).images[0]
+    ecodiff_image = pruned_pipe(
+        prompt=prompt, generator=g_cpu, num_inference_steps=steps
+    ).images[0]
     return original_image, ecodiff_image
 
 
-def on_prune_click(prompt, seed, steps):
+def on_prune_click(prompt, seed, steps, model):
     args = GradioArgs(prompt=prompt, seed=[seed], num_intervention_steps=steps)
-    pipe, pruned_pipe = binary_mask_eval(args)
+    pipe, pruned_pipe = binary_mask_eval(args, model)
     return pipe, pruned_pipe, [("Model Initialized", "green")]
 
 
 def on_generate_click(prompt, seed, steps, pipe, pruned_pipe):
-    original_image, ecodiff_image = generate_images(prompt, seed, steps, pipe, pruned_pipe)
+    original_image, ecodiff_image = generate_images(
+        prompt, seed, steps, pipe, pruned_pipe
+    )
     return original_image, ecodiff_image
 
 
@@ -171,18 +106,23 @@ def create_demo():
         with gr.Row():
             gr.Markdown(
                 """
-                # 🚧 Under Construction 🚧
-                This demo is currently being developed and may not be fully functional. More models and pruning ratios will be supported soon.
-                The current pruned model checkpoint is not optimal and does not provide the best performance.
-                
-                **Note: Please first initialize the model before generating images. This may take up to 5 minutes as it is on CPU.**
+                **Note: Please first initialize the model before generating images. This may take a while to fully loaded.**
                 """
             )
         with gr.Row():
-            model_choice = gr.Dropdown(choices=["SDXL"], value="SDXL", label="Model", scale=1.2)
-            pruning_ratio = gr.Dropdown(choices=["20%"], value="20%", label="Pruning Ratio", scale=1.2)
-            status_label = gr.HighlightedText(label="Model Status", value=[("Model Not Initialized", "red")], scale=1)
-            prune_btn = gr.Button("Initialize Original and Pruned Models", variant="primary", scale=1)
+            model_choice = gr.Radio(
+                choices=["SDXL", "FLUX"], value="SDXL", label="Model", scale=2
+            )
+            pruning_ratio = gr.Text(
+                "20% Pruning Ratio for SDXL, FLUX", label="Pruning Ratio", scale=2
+            )
+            status_label = gr.HighlightedText(
+                label="Model Status", value=[("Model Not Initialized", "red")], scale=1
+            )
+            prune_btn = gr.Button(
+                "Initialize Original and Pruned Models", variant="primary", scale=1
+            )
+
         with gr.Row():
             gr.Markdown(
                 """
@@ -190,9 +130,20 @@ def create_demo():
                 """
             )
         with gr.Row():
-            prompt = gr.Textbox(label="Prompt", value="A clock tower floating in a sea of clouds", scale=3)
+            prompt = gr.Textbox(
+                label="Prompt",
+                value="A clock tower floating in a sea of clouds",
+                scale=3,
+            )
             seed = gr.Number(label="Seed", value=44, precision=0, scale=1)
-            steps = gr.Slider(label="Number of Steps", minimum=1, maximum=100, value=50, step=1, scale=1)
+            steps = gr.Slider(
+                label="Number of Steps",
+                minimum=1,
+                maximum=100,
+                value=50,
+                step=1,
+                scale=1,
+            )
             generate_btn = gr.Button("Generate Images")
         gr.Examples(
             examples=[
@@ -200,6 +151,7 @@ def create_demo():
                 "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
                 "An astronaut riding a green horse",
                 "A delicious ceviche cheesecake slice",
+                "A sprawling cyberpunk metropolis at night, with towering skyscrapers emitting neon lights of every color, holographic billboards advertising alien languages",
             ],
             inputs=[prompt],
         )
@@ -209,6 +161,7 @@ def create_demo():
 
         pipe_state = gr.State(None)
         pruned_pipe_state = gr.State(None)
+
         prompt.submit(
             fn=on_generate_click,
             inputs=[prompt, seed, steps, pipe_state, pruned_pipe_state],
@@ -216,7 +169,7 @@ def create_demo():
         )
         prune_btn.click(
             fn=on_prune_click,
-            inputs=[prompt, seed, steps],
+            inputs=[prompt, seed, steps, model_choice],
             outputs=[pipe_state, pruned_pipe_state, status_label],
         )
         generate_btn.click(
@@ -230,4 +183,4 @@ def create_demo():
 
 if __name__ == "__main__":
     demo = create_demo()
-    demo.launch()
+    demo.launch(share=True)

cross_attn_hook.py

@@ -0,0 +1,632 @@
+import logging
+import os
+from collections import OrderedDict
+from functools import partial
+
+import torch
+
+import re
+
+import math
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from diffusers.models.attention_processor import Attention
+from diffusers.utils import deprecate
+from diffusers.models.embeddings import apply_rotary_emb
+
+
+def scaled_dot_product_attention_atten_weight_only(
+    query, key, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None
+) -> torch.Tensor:
+    L, S = query.size(-2), key.size(-2)
+    scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale
+    attn_bias = torch.zeros(L, S, dtype=query.dtype, device=query.device)
+    if is_causal:
+        assert attn_mask is None
+        temp_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0)
+        attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
+        attn_bias.to(query.dtype)
+
+    if attn_mask is not None:
+        if attn_mask.dtype == torch.bool:
+            attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
+        else:
+            attn_bias += attn_mask
+    attn_weight = query @ key.transpose(-2, -1) * scale_factor
+    attn_weight += attn_bias
+    attn_weight = torch.softmax(attn_weight, dim=-1)
+    attn_weight = torch.dropout(attn_weight, dropout_p, train=True)
+    return attn_weight
+
+
+def apply_rope(xq, xk, freqs_cis):
+    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
+    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
+    xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+    xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+    return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
+
+
+def masking_fn(hidden_states, kwargs):
+    lamb = kwargs["lamb"].view(1, kwargs["lamb"].shape[0], 1, 1)
+    if kwargs.get("masking", None) == "sigmoid":
+        mask = torch.sigmoid(lamb)
+    elif kwargs.get("masking", None) == "binary":
+        mask = lamb
+    elif kwargs.get("masking", None) == "continues2binary":
+        # TODO: this might cause potential issue as it hard threshold at 0
+        mask = (lamb > 0).float()
+    elif kwargs.get("masking", None) == "no_masking":
+        mask = torch.ones_like(lamb)
+    else:
+        raise NotImplementedError
+    epsilon = kwargs.get("epsilon", 0.0)
+    hidden_states = hidden_states * mask + torch.randn_like(hidden_states) * epsilon * (
+        1 - mask
+    )
+    return hidden_states
+
+
+class FluxAttnProcessor2_0_Masking:
+    """Attention processor used typically in processing the SD3-like self-attention projections."""
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError(
+                "FluxAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+            )
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: torch.FloatTensor = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+        *args,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        batch_size, _, _ = (
+            hidden_states.shape
+            if encoder_hidden_states is None
+            else encoder_hidden_states.shape
+        )
+
+        # `sample` projections.
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(hidden_states)
+        value = attn.to_v(hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+        # the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states`
+        if encoder_hidden_states is not None:
+            # `context` projections.
+            encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
+            encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
+            encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
+
+            encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+            encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+            encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+
+            if attn.norm_added_q is not None:
+                encoder_hidden_states_query_proj = attn.norm_added_q(
+                    encoder_hidden_states_query_proj
+                )
+            if attn.norm_added_k is not None:
+                encoder_hidden_states_key_proj = attn.norm_added_k(
+                    encoder_hidden_states_key_proj
+                )
+
+            # attention
+            query = torch.cat([encoder_hidden_states_query_proj, query], dim=2)
+            key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)
+            value = torch.cat([encoder_hidden_states_value_proj, value], dim=2)
+
+        if image_rotary_emb is not None:
+            query = apply_rotary_emb(query, image_rotary_emb)
+            key = apply_rotary_emb(key, image_rotary_emb)
+
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, dropout_p=0.0, is_causal=False
+        )
+
+        if kwargs.get("lamb", None) is not None:
+            hidden_states = masking_fn(hidden_states, kwargs)
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(
+            batch_size, -1, attn.heads * head_dim
+        )
+        hidden_states = hidden_states.to(query.dtype)
+
+        if encoder_hidden_states is not None:
+            encoder_hidden_states, hidden_states = (
+                hidden_states[:, : encoder_hidden_states.shape[1]],
+                hidden_states[:, encoder_hidden_states.shape[1] :],
+            )
+
+            # linear proj
+            hidden_states = attn.to_out[0](hidden_states)
+            # dropout
+            hidden_states = attn.to_out[1](hidden_states)
+            encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+            return hidden_states, encoder_hidden_states
+        else:
+            return hidden_states
+
+
+class AttnProcessor2_0_Masking:
+    r"""
+    Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
+    """
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError(
+                "AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+            )
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        temb: Optional[torch.Tensor] = None,
+        *args,
+        **kwargs,
+    ):
+        if len(args) > 0 or kwargs.get("scale", None) is not None:
+            deprecation_message = (
+                "The `scale` argument is deprecated and will be ignored. "
+                "Please remove it, as passing it will raise an error "
+                "in the future. `scale` should directly be passed while "
+                "calling the underlying pipeline component i.e., via "
+                "`cross_attention_kwargs`."
+            )
+            deprecate("scale", "1.0.0", deprecation_message)
+
+        residual = hidden_states
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(
+                batch_size, channel, height * width
+            ).transpose(1, 2)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape
+            if encoder_hidden_states is None
+            else encoder_hidden_states.shape
+        )
+
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(
+                attention_mask, sequence_length, batch_size
+            )
+            # scaled_dot_product_attention expects attention_mask shape to be
+            # (batch, heads, source_length, target_length)
+            attention_mask = attention_mask.view(
+                batch_size, attn.heads, -1, attention_mask.shape[-1]
+            )
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(
+                1, 2
+            )
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(
+                encoder_hidden_states
+            )
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        if getattr(attn, "norm_q", None) is not None:
+            query = attn.norm_q(query)
+
+        if getattr(attn, "norm_k", None) is not None:
+            key = attn.norm_k(key)
+
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        if kwargs.get("return_attention", True):
+            # add the attention output from F.scaled_dot_product_attention
+            attn_weight = scaled_dot_product_attention_atten_weight_only(
+                query, key, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+            )
+            hidden_states_aft_attention_ops = hidden_states.clone()
+            attn_weight_old = attn_weight.to(hidden_states.device).clone()
+        else:
+            hidden_states_aft_attention_ops = None
+            attn_weight_old = None
+
+        # masking for the hidden_states after the attention ops
+        if kwargs.get("lamb", None) is not None:
+            hidden_states = masking_fn(hidden_states, kwargs)
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(
+            batch_size, -1, attn.heads * head_dim
+        )
+        hidden_states = hidden_states.to(query.dtype)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(
+                batch_size, channel, height, width
+            )
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states, hidden_states_aft_attention_ops, attn_weight_old
+
+
+class BaseCrossAttentionHooker:
+    def __init__(
+        self,
+        pipeline,
+        regex,
+        dtype,
+        head_num_filter,
+        masking,
+        model_name,
+        attn_name,
+        use_log,
+        eps,
+    ):
+        self.pipeline = pipeline
+        # unet for SD2 SDXL, transformer for SD3, FLUX DIT
+        self.net = pipeline.unet if hasattr(pipeline, "unet") else pipeline.transformer
+        self.model_name = model_name
+        self.module_heads = OrderedDict()
+        self.masking = masking
+        self.hook_dict = {}
+        self.regex = regex
+        self.dtype = dtype
+        self.head_num_filter = head_num_filter
+        self.attn_name = attn_name
+        self.logger = logging.getLogger(__name__)
+        self.use_log = use_log  # use log parameter to control hard_discrete
+        self.eps = eps
+
+    def add_hooks_to_cross_attention(self, hook_fn: callable):
+        """
+        Add forward hooks to every cross attention
+        :param hook_fn: a callable to be added to torch nn module as a hook
+        :return:
+        """
+        total_hooks = 0
+        for name, module in self.net.named_modules():
+            name_last_word = name.split(".")[-1]
+            if self.attn_name in name_last_word:
+                if re.match(self.regex, name):
+                    hook_fn = partial(hook_fn, name=name)
+                    hook = module.register_forward_hook(hook_fn, with_kwargs=True)
+                    self.hook_dict[name] = hook
+                    self.module_heads[name] = module.heads
+                    self.logger.info(
+                        f"Adding hook to {name}, module.heads: {module.heads}"
+                    )
+                    total_hooks += 1
+        self.logger.info(f"Total hooks added: {total_hooks}")
+
+    def clear_hooks(self):
+        """clear all hooks"""
+        for hook in self.hook_dict.values():
+            hook.remove()
+        self.hook_dict.clear()
+
+
+class CrossAttentionExtractionHook(BaseCrossAttentionHooker):
+    def __init__(
+        self,
+        pipeline,
+        dtype,
+        head_num_filter,
+        masking,
+        dst,
+        regex=None,
+        epsilon=0.0,
+        binary=False,
+        return_attention=False,
+        model_name="sdxl",
+        attn_name="attn",
+        use_log=False,
+        eps=1e-6,
+    ):
+        super().__init__(
+            pipeline,
+            regex,
+            dtype,
+            head_num_filter,
+            masking=masking,
+            model_name=model_name,
+            attn_name=attn_name,
+            use_log=use_log,
+            eps=eps,
+        )
+        if model_name == "sdxl":
+            self.attention_processor = AttnProcessor2_0_Masking()
+        elif model_name == "flux":
+            self.attention_processor = FluxAttnProcessor2_0_Masking()
+        self.lambs = []
+        self.lambs_module_names = []
+        self.cross_attn = []
+        self.hook_counter = 0
+        self.device = (
+            self.pipeline.unet.device
+            if hasattr(self.pipeline, "unet")
+            else self.pipeline.transformer.device
+        )
+        self.dst = dst
+        self.epsilon = epsilon
+        self.binary = binary
+        self.return_attention = return_attention
+        self.model_name = model_name
+
+    def clean_cross_attn(self):
+        self.cross_attn = []
+
+    def validate_dst(self):
+        if os.path.exists(self.dst):
+            raise ValueError(f"Destination {self.dst} already exists")
+
+    def save(self, name: str = None):
+        if name is not None:
+            dst = os.path.join(os.path.dirname(self.dst), name)
+        else:
+            dst = self.dst
+        dst_dir = os.path.dirname(dst)
+        if not os.path.exists(dst_dir):
+            self.logger.info(f"Creating directory {dst_dir}")
+            os.makedirs(dst_dir)
+        torch.save(self.lambs, dst)
+
+    @property
+    def get_lambda_block_names(self):
+        return self.lambs_module_names
+
+    def load(self, device, threshold=2.5):
+        if os.path.exists(self.dst):
+            self.logger.info(f"loading lambda from {self.dst}")
+            self.lambs = torch.load(self.dst, weights_only=True, map_location=device)
+            if self.binary:
+                # set binary masking for each lambda by using clamp
+                self.lambs = [
+                    (torch.relu(lamb - threshold) > 0).float() for lamb in self.lambs
+                ]
+        else:
+            self.logger.info("skipping loading, training from scratch")
+
+    def binarize(self, scope: str, ratio: float):
+        assert scope in ["local", "global"], "scope must be either local or global"
+        assert (
+            not self.binary
+        ), "binarization is not supported when using binary mask already"
+        if scope == "local":
+            # Local binarization
+            for i, lamb in enumerate(self.lambs):
+                num_heads = lamb.size(0)
+                num_activate_heads = int(num_heads * ratio)
+                # Sort the lambda values with stable sorting to maintain order for equal values
+                sorted_lamb, sorted_indices = torch.sort(
+                    lamb, descending=True, stable=True
+                )
+                # Find the threshold value
+                threshold = sorted_lamb[num_activate_heads - 1]
+                # Create a mask based on the sorted indices
+                mask = torch.zeros_like(lamb)
+                mask[sorted_indices[:num_activate_heads]] = 1.0
+                # Binarize the lambda based on the threshold and the mask
+                self.lambs[i] = torch.where(
+                    lamb > threshold, torch.ones_like(lamb), mask
+                )
+        else:
+            # Global binarization
+            all_lambs = torch.cat([lamb.flatten() for lamb in self.lambs])
+            num_total = all_lambs.numel()
+            num_activate = int(num_total * ratio)
+            # Sort all lambda values globally with stable sorting
+            sorted_lambs, sorted_indices = torch.sort(
+                all_lambs, descending=True, stable=True
+            )
+            # Find the global threshold value
+            threshold = sorted_lambs[num_activate - 1]
+            # Create a global mask based on the sorted indices
+            global_mask = torch.zeros_like(all_lambs)
+            global_mask[sorted_indices[:num_activate]] = 1.0
+            # Binarize all lambdas based on the global threshold and mask
+            start_idx = 0
+            for i in range(len(self.lambs)):
+                end_idx = start_idx + self.lambs[i].numel()
+                lamb_mask = global_mask[start_idx:end_idx].reshape(self.lambs[i].shape)
+                self.lambs[i] = torch.where(
+                    self.lambs[i] > threshold, torch.ones_like(self.lambs[i]), lamb_mask
+                )
+                start_idx = end_idx
+        self.binary = True
+
+    def bizarize_threshold(self, threshold: float):
+        """
+        Binarize lambda values based on a predefined threshold.
+        :param threshold: The threshold value for binarization
+        """
+        assert (
+            not self.binary
+        ), "Binarization is not supported when using binary mask already"
+
+        for i in range(len(self.lambs)):
+            self.lambs[i] = (self.lambs[i] >= threshold).float()
+
+        self.binary = True
+
+    def get_cross_attn_extraction_hook(self, init_value=1.0):
+        """get a hook function to extract cross attention"""
+
+        # the reason to use a function inside a function is to save the extracted cross attention
+        def hook_fn(module, args, kwargs, output, name):
+            # initialize lambda with acual head dim in the first run
+            if self.lambs[self.hook_counter] is None:
+                self.lambs[self.hook_counter] = (
+                    torch.ones(
+                        module.heads, device=self.pipeline.device, dtype=self.dtype
+                    )
+                    * init_value
+                )
+                # Only set requires_grad to True when the head number is larger than the filter
+                if self.head_num_filter <= module.heads:
+                    self.lambs[self.hook_counter].requires_grad = True
+
+                # load attn lambda module name for logging
+                self.lambs_module_names[self.hook_counter] = name
+
+            if self.model_name == "sdxl":
+                hidden_states, _, attention_output = self.attention_processor(
+                    module,
+                    args[0],
+                    encoder_hidden_states=kwargs["encoder_hidden_states"],
+                    attention_mask=kwargs["attention_mask"],
+                    lamb=self.lambs[self.hook_counter],
+                    masking=self.masking,
+                    epsilon=self.epsilon,
+                    return_attention=self.return_attention,
+                    use_log=self.use_log,
+                    eps=self.eps,
+                )
+                if attention_output is not None:
+                    self.cross_attn.append(attention_output)
+                self.hook_counter += 1
+                self.hook_counter %= len(self.lambs)
+                return hidden_states
+            elif self.model_name == "flux":
+                encoder_hidden_states = kwargs.get("encoder_hidden_states", None)
+                # flux has two different attention processors, FluxSingleAttnProcessor and FluxAttnProcessor
+                if "single" in name:
+                    hidden_states = self.attention_processor(
+                        module,
+                        hidden_states=kwargs.get("hidden_states", None),
+                        encoder_hidden_states=encoder_hidden_states,
+                        attention_mask=kwargs.get("attention_mask", None),
+                        image_rotary_emb=kwargs.get("image_rotary_emb", None),
+                        lamb=self.lambs[self.hook_counter],
+                        masking=self.masking,
+                        epsilon=self.epsilon,
+                        use_log=self.use_log,
+                        eps=self.eps,
+                    )
+                    self.hook_counter += 1
+                    self.hook_counter %= len(self.lambs)
+                    return hidden_states
+                else:
+                    hidden_states, encoder_hidden_states = self.attention_processor(
+                        module,
+                        hidden_states=kwargs.get("hidden_states", None),
+                        encoder_hidden_states=encoder_hidden_states,
+                        attention_mask=kwargs.get("attention_mask", None),
+                        image_rotary_emb=kwargs.get("image_rotary_emb", None),
+                        lamb=self.lambs[self.hook_counter],
+                        masking=self.masking,
+                        epsilon=self.epsilon,
+                        use_log=self.use_log,
+                        eps=self.eps,
+                    )
+                    self.hook_counter += 1
+                    self.hook_counter %= len(self.lambs)
+                    return hidden_states, encoder_hidden_states
+
+        return hook_fn
+
+    def add_hooks(self, init_value=1.0):
+        hook_fn = self.get_cross_attn_extraction_hook(init_value)
+        self.add_hooks_to_cross_attention(hook_fn)
+        # initialize the lambda
+        self.lambs = [None] * len(self.module_heads)
+        # initialize the lambda module names
+        self.lambs_module_names = [None] * len(self.module_heads)
+
+    def get_process_cross_attn_result(self, text_seq_length, timestep: int = -1):
+        if isinstance(timestep, str):
+            timestep = int(timestep)
+        # num_lambda_block contains lambda (head masking)
+        num_lambda_block = len(self.lambs)
+
+        # get the start and end position of the timestep
+        start_pos = timestep * num_lambda_block
+        end_pos = (timestep + 1) * num_lambda_block
+        if end_pos > len(self.cross_attn):
+            raise ValueError(f"timestep {timestep} is out of range")
+
+        # list[cross_attn_map] num_layer x [batch, num_heads, seq_vis_tokens, seq_text_tokens]
+        attn_maps = self.cross_attn[start_pos:end_pos]
+
+        def heatmap(attn_list, attn_idx, head_idx, text_idx):
+            # only select second element in the tuple (with text guided attention)
+            # layer_idx, 1, head_idx, seq_vis_tokens, seq_text_tokens
+            map = attn_list[attn_idx][1][head_idx][:][:, text_idx]
+            # get the size of the heatmap
+            size = int(map.shape[0] ** 0.5)
+            map = map.view(size, size, 1)
+            data = map.cpu().float().numpy()
+            return data
+
+        output_dict = {}
+        for lambda_block_idx, lambda_block_name in zip(
+            range(num_lambda_block), self.lambs_module_names
+        ):
+            data_list = []
+            for head_idx in range(len(self.lambs[lambda_block_idx])):
+                for token_idx in range(text_seq_length):
+                    # number of heatmap is equal to the number of tokens in the text sequence X number of heads
+                    data_list.append(
+                        heatmap(attn_maps, lambda_block_idx, head_idx, token_idx)
+                    )
+            output_dict[lambda_block_name] = {
+                "attn_map": data_list,
+                "lambda": self.lambs[lambda_block_idx],
+            }
+        return output_dict

ffn_hooker.py

@@ -0,0 +1,224 @@
+import logging
+import os
+from collections import OrderedDict
+from functools import partial
+
+import diffusers
+import torch
+from torch import nn
+
+import re
+
+
+class FeedForwardHooker:
+    def __init__(
+        self,
+        pipeline: nn.Module,
+        regex: str,
+        dtype: torch.dtype,
+        masking: str,
+        dst: str,
+        epsilon: float = 0.0,
+        eps: float = 1e-6,
+        use_log: bool = False,
+        binary: bool = False,
+    ):
+        self.pipeline = pipeline
+        self.net = pipeline.unet if hasattr(pipeline, "unet") else pipeline.transformer
+        self.logger = logging.getLogger(__name__)
+        self.dtype = dtype
+        self.regex = regex
+        self.hook_dict = {}
+        self.masking = masking
+        self.dst = dst
+        self.epsilon = epsilon
+        self.eps = eps
+        self.use_log = use_log
+        self.lambs = []
+        self.lambs_module_names = []  # store the module names for each lambda block
+        self.hook_counter = 0
+        self.module_neurons = OrderedDict()
+        self.binary = (
+            binary  # default, need to discuss if we need to keep this attribute or not
+        )
+
+    def add_hooks_to_ff(self, hook_fn: callable):
+        total_hooks = 0
+        for name, module in self.net.named_modules():
+            name_last_word = name.split(".")[-1]
+            if "ff" in name_last_word:
+                if re.match(self.regex, name):
+                    hook_fn_with_name = partial(hook_fn, name=name)
+                    actual_module = module.net[0]
+                    hook = actual_module.register_forward_hook(
+                        hook_fn_with_name, with_kwargs=True
+                    )
+                    self.hook_dict[name] = hook
+
+                    if isinstance(
+                        actual_module, diffusers.models.activations.GEGLU
+                    ):  # geglu
+                        # due to the GEGLU chunking, we need to divide by 2
+                        self.module_neurons[name] = actual_module.proj.out_features // 2
+                    elif isinstance(
+                        actual_module, diffusers.models.activations.GELU
+                    ):  # gelu
+                        self.module_neurons[name] = actual_module.proj.out_features
+                    else:
+                        raise NotImplementedError(
+                            f"Module {name} is not implemented, please check"
+                        )
+                    self.logger.info(
+                        f"Adding hook to {name}, neurons: {self.module_neurons[name]}"
+                    )
+                    total_hooks += 1
+        self.logger.info(f"Total hooks added: {total_hooks}")
+        return self.hook_dict
+
+    def add_hooks(self, init_value=1.0):
+        hook_fn = self.get_ff_masking_hook(init_value)
+        self.add_hooks_to_ff(hook_fn)
+        # initialize the lambda
+        self.lambs = [None] * len(self.hook_dict)
+        # initialize the lambda module names
+        self.lambs_module_names = [None] * len(self.hook_dict)
+
+    def clear_hooks(self):
+        """clear all hooks"""
+        for hook in self.hook_dict.values():
+            hook.remove()
+        self.hook_dict.clear()
+
+    def save(self, name: str = None):
+        if name is not None:
+            dst = os.path.join(os.path.dirname(self.dst), name)
+        else:
+            dst = self.dst
+        dst_dir = os.path.dirname(dst)
+        if not os.path.exists(dst_dir):
+            self.logger.info(f"Creating directory {dst_dir}")
+            os.makedirs(dst_dir)
+        torch.save(self.lambs, dst)
+
+    @property
+    def get_lambda_block_names(self):
+        return self.lambs_module_names
+
+    def load(self, device, threshold=2.5):
+        if os.path.exists(self.dst):
+            self.logger.info(f"loading lambda from {self.dst}")
+            self.lambs = torch.load(self.dst, weights_only=True, map_location=device)
+            if self.binary:
+                # set binary masking for each lambda by using clamp
+                self.lambs = [
+                    (torch.relu(lamb - threshold) > 0).float() for lamb in self.lambs
+                ]
+            else:
+                self.lambs = [torch.clamp(lamb, min=0.0) for lamb in self.lambs]
+            # self.lambs_module_names = [None for _ in self.lambs]
+        else:
+            self.logger.info("skipping loading, training from scratch")
+
+    def binarize(self, scope: str, ratio: float):
+        assert scope in ["local", "global"], "scope must be either local or global"
+        assert (
+            not self.binary
+        ), "binarization is not supported when using binary mask already"
+        if scope == "local":
+            # Local binarization
+            for i, lamb in enumerate(self.lambs):
+                num_heads = lamb.size(0)
+                num_activate_heads = int(num_heads * ratio)
+                # Sort the lambda values with stable sorting to maintain order for equal values
+                sorted_lamb, sorted_indices = torch.sort(
+                    lamb, descending=True, stable=True
+                )
+                # Find the threshold value
+                threshold = sorted_lamb[num_activate_heads - 1]
+                # Create a mask based on the sorted indices
+                mask = torch.zeros_like(lamb)
+                mask[sorted_indices[:num_activate_heads]] = 1.0
+                # Binarize the lambda based on the threshold and the mask
+                self.lambs[i] = torch.where(
+                    lamb > threshold, torch.ones_like(lamb), mask
+                )
+        else:
+            # Global binarization
+            all_lambs = torch.cat([lamb.flatten() for lamb in self.lambs])
+            num_total = all_lambs.numel()
+            num_activate = int(num_total * ratio)
+            # Sort all lambda values globally with stable sorting
+            sorted_lambs, sorted_indices = torch.sort(
+                all_lambs, descending=True, stable=True
+            )
+            # Find the global threshold value
+            threshold = sorted_lambs[num_activate - 1]
+            # Create a global mask based on the sorted indices
+            global_mask = torch.zeros_like(all_lambs)
+            global_mask[sorted_indices[:num_activate]] = 1.0
+            # Binarize all lambdas based on the global threshold and mask
+            start_idx = 0
+            for i in range(len(self.lambs)):
+                end_idx = start_idx + self.lambs[i].numel()
+                lamb_mask = global_mask[start_idx:end_idx].reshape(self.lambs[i].shape)
+                self.lambs[i] = torch.where(
+                    self.lambs[i] > threshold, torch.ones_like(self.lambs[i]), lamb_mask
+                )
+                start_idx = end_idx
+        self.binary = True
+
+    @staticmethod
+    def masking_fn(hidden_states, **kwargs):
+        hidden_states_dtype = hidden_states.dtype
+        lamb = kwargs["lamb"].view(1, 1, kwargs["lamb"].shape[0])
+        if kwargs.get("masking", None) == "sigmoid":
+            mask = torch.sigmoid(lamb)
+        elif kwargs.get("masking", None) == "binary":
+            mask = lamb
+        elif kwargs.get("masking", None) == "continues2binary":
+            # TODO: this might cause potential issue as it hard threshold at 0
+            mask = (lamb > 0).float()
+        elif kwargs.get("masking", None) == "no_masking":
+            mask = torch.ones_like(lamb)
+        else:
+            raise NotImplementedError
+        epsilon = kwargs.get("epsilon", 0.0)
+        hidden_states = hidden_states * mask + torch.randn_like(
+            hidden_states
+        ) * epsilon * (1 - mask)
+        return hidden_states.to(hidden_states_dtype)
+
+    def get_ff_masking_hook(self, init_value=1.0):
+        """
+        Get a hook function to mask feed forward layer
+        """
+
+        def hook_fn(module, args, kwargs, output, name):
+            # initialize lambda with acual head dim in the first run
+            if self.lambs[self.hook_counter] is None:
+                self.lambs[self.hook_counter] = (
+                    torch.ones(
+                        self.module_neurons[name],
+                        device=self.pipeline.device,
+                        dtype=self.dtype,
+                    )
+                    * init_value
+                )
+                self.lambs[self.hook_counter].requires_grad = True
+                # load ff lambda module name for logging
+                self.lambs_module_names[self.hook_counter] = name
+
+            # perform masking
+            output = self.masking_fn(
+                output,
+                masking=self.masking,
+                lamb=self.lambs[self.hook_counter],
+                epsilon=self.epsilon,
+                eps=self.eps,
+                use_log=self.use_log,
+            )
+            self.hook_counter += 1
+            self.hook_counter %= len(self.lambs)
+            return output
+
+        return hook_fn

norm_attn_hook.py

@@ -0,0 +1,242 @@
+# TODO should be a parent class for all the hooks !! for the offical repo
+# 1: FLUX Norm
+
+import logging
+import os
+from collections import OrderedDict
+from functools import partial
+
+import torch
+from torch import nn
+
+import re
+
+
+class NormHooker:
+    def __init__(
+        self,
+        pipeline: nn.Module,
+        regex: str,
+        dtype: torch.dtype,
+        masking: str,
+        dst: str,
+        epsilon: float = 0.0,
+        eps: float = 1e-6,
+        use_log: bool = False,
+        binary: bool = False,
+    ):
+        self.pipeline = pipeline
+        self.net = pipeline.unet if hasattr(pipeline, "unet") else pipeline.transformer
+        self.logger = logging.getLogger(__name__)
+        self.dtype = dtype
+        self.regex = regex
+        self.hook_dict = {}
+        self.masking = masking
+        self.dst = dst
+        self.epsilon = epsilon
+        self.eps = eps
+        self.use_log = use_log
+        self.lambs = []
+        self.lambs_module_names = []  # store the module names for each lambda block
+        self.hook_counter = 0
+        self.module_neurons = OrderedDict()
+        self.binary = (
+            binary  # default, need to discuss if we need to keep this attribute or not
+        )
+
+    def add_hooks_to_norm(self, hook_fn: callable):
+        """
+        Add forward hooks to every feed forward layer matching the regex
+        :param hook_fn: a callable to be added to torch nn module as a hook
+        :return: dictionary of added hooks
+        """
+        total_hooks = 0
+        for name, module in self.net.named_modules():
+            name_last_word = name.split(".")[-1]
+            if "norm1_context" in name_last_word:
+                if re.match(self.regex, name):
+                    hook_fn_with_name = partial(hook_fn, name=name)
+
+                    if hasattr(module, "linear"):
+                        actual_module = module.linear
+                    else:
+                        if isinstance(module, nn.Linear):
+                            actual_module = module
+                        else:
+                            continue
+
+                    hook = actual_module.register_forward_hook(
+                        hook_fn_with_name, with_kwargs=True
+                    )
+                    self.hook_dict[name] = hook
+
+                    # AdaLayerNormZero
+                    if isinstance(actual_module, torch.nn.Linear):
+                        self.module_neurons[name] = actual_module.out_features
+                    else:
+                        raise NotImplementedError(
+                            f"Module {name} is not implemented, please check"
+                        )
+                    self.logger.info(
+                        f"Adding hook to {name}, neurons: {self.module_neurons[name]}"
+                    )
+                    total_hooks += 1
+        self.logger.info(f"Total hooks added: {total_hooks}")
+        return self.hook_dict
+
+    def add_hooks(self, init_value=1.0):
+        hook_fn = self.get_norm_masking_hook(init_value)
+        self.add_hooks_to_norm(hook_fn)
+        # initialize the lambda
+        self.lambs = [None] * len(self.hook_dict)
+        # initialize the lambda module names
+        self.lambs_module_names = [None] * len(self.hook_dict)
+
+    def clear_hooks(self):
+        """clear all hooks"""
+        for hook in self.hook_dict.values():
+            hook.remove()
+        self.hook_dict.clear()
+
+    def save(self, name: str = None):
+        if name is not None:
+            dst = os.path.join(os.path.dirname(self.dst), name)
+        else:
+            dst = self.dst
+        dst_dir = os.path.dirname(dst)
+        if not os.path.exists(dst_dir):
+            self.logger.info(f"Creating directory {dst_dir}")
+            os.makedirs(dst_dir)
+        torch.save(self.lambs, dst)
+
+    @property
+    def get_lambda_block_names(self):
+        return self.lambs_module_names
+
+    def load(self, device, threshold):
+        if os.path.exists(self.dst):
+            self.logger.info(f"loading lambda from {self.dst}")
+            self.lambs = torch.load(self.dst, weights_only=True, map_location=device)
+            if self.binary:
+                # set binary masking for each lambda by using clamp
+                self.lambs = [
+                    (torch.relu(lamb - threshold) > 0).float() for lamb in self.lambs
+                ]
+            else:
+                self.lambs = [torch.clamp(lamb, min=0.0) for lamb in self.lambs]
+            # self.lambs_module_names = [None for _ in self.lambs]
+        else:
+            self.logger.info("skipping loading, training from scratch")
+
+    def binarize(self, scope: str, ratio: float):
+        """
+        binarize lambda to be 0 or 1
+        :param scope: either locally (sparsity within layer) or globally (sparsity within model)
+        :param ratio: the ratio of the number of 1s to the total number of elements
+        """
+        assert scope in ["local", "global"], "scope must be either local or global"
+        assert (
+            not self.binary
+        ), "binarization is not supported when using binary mask already"
+        if scope == "local":
+            # Local binarization
+            for i, lamb in enumerate(self.lambs):
+                num_heads = lamb.size(0)
+                num_activate_heads = int(num_heads * ratio)
+                # Sort the lambda values with stable sorting to maintain order for equal values
+                sorted_lamb, sorted_indices = torch.sort(
+                    lamb, descending=True, stable=True
+                )
+                # Find the threshold value
+                threshold = sorted_lamb[num_activate_heads - 1]
+                # Create a mask based on the sorted indices
+                mask = torch.zeros_like(lamb)
+                mask[sorted_indices[:num_activate_heads]] = 1.0
+                # Binarize the lambda based on the threshold and the mask
+                self.lambs[i] = torch.where(
+                    lamb > threshold, torch.ones_like(lamb), mask
+                )
+        else:
+            # Global binarization
+            all_lambs = torch.cat([lamb.flatten() for lamb in self.lambs])
+            num_total = all_lambs.numel()
+            num_activate = int(num_total * ratio)
+            # Sort all lambda values globally with stable sorting
+            sorted_lambs, sorted_indices = torch.sort(
+                all_lambs, descending=True, stable=True
+            )
+            # Find the global threshold value
+            threshold = sorted_lambs[num_activate - 1]
+            # Create a global mask based on the sorted indices
+            global_mask = torch.zeros_like(all_lambs)
+            global_mask[sorted_indices[:num_activate]] = 1.0
+            # Binarize all lambdas based on the global threshold and mask
+            start_idx = 0
+            for i in range(len(self.lambs)):
+                end_idx = start_idx + self.lambs[i].numel()
+                lamb_mask = global_mask[start_idx:end_idx].reshape(self.lambs[i].shape)
+                self.lambs[i] = torch.where(
+                    self.lambs[i] > threshold, torch.ones_like(self.lambs[i]), lamb_mask
+                )
+                start_idx = end_idx
+        self.binary = True
+
+    @staticmethod
+    def masking_fn(hidden_states, **kwargs):
+        hidden_states_dtype = hidden_states.dtype
+        lamb = kwargs["lamb"].view(1, 1, kwargs["lamb"].shape[0])
+        if kwargs.get("masking", None) == "sigmoid":
+            mask = torch.sigmoid(lamb)
+        elif kwargs.get("masking", None) == "binary":
+            mask = lamb
+        elif kwargs.get("masking", None) == "continues2binary":
+            # TODO: this might cause potential issue as it hard threshold at 0
+            mask = (lamb > 0).float()
+        elif kwargs.get("masking", None) == "no_masking":
+            mask = torch.ones_like(lamb)
+        else:
+            raise NotImplementedError
+        epsilon = kwargs.get("epsilon", 0.0)
+
+        if hidden_states.dim() == 2:
+            mask = mask.squeeze(1)
+
+        hidden_states = hidden_states * mask + torch.randn_like(
+            hidden_states
+        ) * epsilon * (1 - mask)
+        return hidden_states.to(hidden_states_dtype)
+
+    def get_norm_masking_hook(self, init_value=1.0):
+        """
+        Get a hook function to mask feed forward layer
+        """
+
+        def hook_fn(module, args, kwargs, output, name):
+            # initialize lambda with acual head dim in the first run
+            if self.lambs[self.hook_counter] is None:
+                self.lambs[self.hook_counter] = (
+                    torch.ones(
+                        self.module_neurons[name],
+                        device=self.pipeline.device,
+                        dtype=self.dtype,
+                    )
+                    * init_value
+                )
+                self.lambs[self.hook_counter].requires_grad = True
+                # load norm lambda module name for logging
+                self.lambs_module_names[self.hook_counter] = name
+
+            # perform masking
+            output = self.masking_fn(
+                output,
+                masking=self.masking,
+                lamb=self.lambs[self.hook_counter],
+                epsilon=self.epsilon,
+                eps=self.eps,
+                use_log=self.use_log,
+            )
+            self.hook_counter += 1
+            self.hook_counter %= len(self.lambs)
+            return output
+
+        return hook_fn

utils.py

@@ -0,0 +1,310 @@
+import os
+import pickle
+from copy import deepcopy
+from typing import Optional
+
+import torch
+from diffusers.models.activations import GEGLU, GELU
+from cross_attn_hook import CrossAttentionExtractionHook
+from ffn_hooker import FeedForwardHooker
+from norm_attn_hook import NormHooker
+
+
+# create dummy module for skip connection
+class SkipConnection(torch.nn.Module):
+    def __init__(self):
+        super(SkipConnection, self).__init__()
+
+    def forward(*args, **kwargs):
+        return args[1]
+
+
+def calculate_mask_sparsity(hooker, threshold: Optional[float] = None):
+    total_num_lambs = 0
+    num_activate_lambs = 0
+    binary = getattr(
+        hooker, "binary", None
+    )  # if binary is not present, it will return None for ff_hooks
+    for lamb in hooker.lambs:
+        total_num_lambs += lamb.size(0)
+        if binary:
+            assert threshold is None, "threshold should be None for binary mask"
+            num_activate_lambs += lamb.sum().item()
+        else:
+            assert (
+                threshold is not None
+            ), "threshold must be provided for non-binary mask"
+            num_activate_lambs += (lamb >= threshold).sum().item()
+    return total_num_lambs, num_activate_lambs, num_activate_lambs / total_num_lambs
+
+
+def create_pipeline(
+    pipe,
+    model_id,
+    device,
+    torch_dtype,
+    save_pt=None,
+    lambda_threshold: float = 1,
+    binary=True,
+    epsilon=0.0,
+    masking="binary",
+    attn_name="attn",
+    return_hooker=False,
+    scope=None,
+    ratio=None,
+):
+    """
+    create the pipeline and optionally load the saved mask
+    """
+    pipe.to(device)
+    pipe.vae.requires_grad_(False)
+    if hasattr(pipe, "unet"):
+        pipe.unet.requires_grad_(False)
+    else:
+        pipe.transformer.requires_grad_(False)
+    if save_pt:
+        # TODO should merge all the hooks checkpoint into one
+        if "ff.pt" in save_pt or "attn.pt" in save_pt:
+            save_pts = get_save_pts(save_pt)
+
+            cross_attn_hooker = CrossAttentionExtractionHook(
+                pipe,
+                model_name=model_id,
+                regex=".*",
+                dtype=torch_dtype,
+                head_num_filter=1,
+                masking=masking,  # need to change to binary during inference
+                dst=save_pts["attn"],
+                epsilon=epsilon,
+                attn_name=attn_name,
+                binary=binary,
+            )
+            cross_attn_hooker.add_hooks(init_value=1)
+
+            ff_hooker = FeedForwardHooker(
+                pipe,
+                regex=".*",
+                dtype=torch_dtype,
+                masking=masking,
+                dst=save_pts["ff"],
+                epsilon=epsilon,
+                binary=binary,
+            )
+            ff_hooker.add_hooks(init_value=1)
+
+            if os.path.exists(save_pts["norm"]):
+                norm_hooker = NormHooker(
+                    pipe,
+                    regex=".*",
+                    dtype=torch_dtype,
+                    masking=masking,
+                    dst=save_pts["norm"],
+                    epsilon=epsilon,
+                    binary=binary,
+                )
+                norm_hooker.add_hooks(init_value=1)
+            else:
+                norm_hooker = None
+
+            _ = pipe("abc", num_inference_steps=1)
+            cross_attn_hooker.load(device=device, threshold=lambda_threshold)
+            ff_hooker.load(device=device, threshold=lambda_threshold)
+            if norm_hooker:
+                norm_hooker.load(device=device, threshold=lambda_threshold)
+            if scope == "local" or scope == "global":
+                if isinstance(ratio, float):
+                    attn_hooker_ratio = ratio
+                    ff_hooker_ratio = ratio
+                else:
+                    attn_hooker_ratio, ff_hooker_ratio = ratio[0], ratio[1]
+
+                    if norm_hooker:
+                        if len(ratio) < 3:
+                            raise ValueError("Need to provide ratio for norm layer")
+                        norm_hooker_ratio = ratio[2]
+
+                cross_attn_hooker.binarize(scope, attn_hooker_ratio)
+                ff_hooker.binarize(scope, ff_hooker_ratio)
+                if norm_hooker:
+                    norm_hooker.binarize(scope, norm_hooker_ratio)
+            hookers = [cross_attn_hooker, ff_hooker]
+            if norm_hooker:
+                hookers.append(norm_hooker)
+
+    if return_hooker:
+        return pipe, hookers
+    else:
+        return pipe
+
+
+def linear_layer_pruning(module, lamb):
+    heads_to_keep = torch.nonzero(lamb).squeeze()
+    if len(heads_to_keep.shape) == 0:
+        # if only one head is kept, or none
+        heads_to_keep = heads_to_keep.unsqueeze(0)
+
+    modules_to_remove = [module.to_k, module.to_q, module.to_v]
+    new_heads = int(lamb.sum().item())
+
+    if new_heads == 0:
+        return SkipConnection()
+
+    for module_to_remove in modules_to_remove:
+        # get head dimension
+        inner_dim = module_to_remove.out_features // module.heads
+        # place holder for the rows to keep
+        rows_to_keep = torch.zeros(
+            module_to_remove.out_features,
+            dtype=torch.bool,
+            device=module_to_remove.weight.device,
+        )
+
+        for idx in heads_to_keep:
+            rows_to_keep[idx * inner_dim : (idx + 1) * inner_dim] = True
+
+        # overwrite the inner projection with masked projection
+        module_to_remove.weight.data = module_to_remove.weight.data[rows_to_keep, :]
+        if module_to_remove.bias is not None:
+            module_to_remove.bias.data = module_to_remove.bias.data[rows_to_keep]
+        module_to_remove.out_features = int(sum(rows_to_keep).item())
+
+    # Also update the output projection layer if available, (for FLUXSingleAttnProcessor2_0)
+    # with column masking, dim 1
+    if getattr(module, "to_out", None) is not None:
+        module.to_out[0].weight.data = module.to_out[0].weight.data[:, rows_to_keep]
+        module.to_out[0].in_features = int(sum(rows_to_keep).item())
+
+    # update parameters in the attention module
+    module.inner_dim = module.inner_dim // module.heads * new_heads
+    try:
+        module.query_dim = module.query_dim // module.heads * new_heads
+        module.inner_kv_dim = module.inner_kv_dim // module.heads * new_heads
+    except:
+        pass
+    module.cross_attention_dim = module.cross_attention_dim // module.heads * new_heads
+    module.heads = new_heads
+    return module
+
+
+def ffn_linear_layer_pruning(module, lamb):
+    lambda_to_keep = torch.nonzero(lamb).squeeze()
+    if len(lambda_to_keep) == 0:
+        return SkipConnection()
+
+    num_lambda = len(lambda_to_keep)
+
+    if isinstance(module.net[0], GELU):
+        # linear layer weight remove before activation
+        module.net[0].proj.weight.data = module.net[0].proj.weight.data[
+            lambda_to_keep, :
+        ]
+        module.net[0].proj.out_features = num_lambda
+        if module.net[0].proj.bias is not None:
+            module.net[0].proj.bias.data = module.net[0].proj.bias.data[lambda_to_keep]
+
+        update_act = GELU(module.net[0].proj.in_features, num_lambda)
+        update_act.proj = module.net[0].proj
+        module.net[0] = update_act
+    elif isinstance(module.net[0], GEGLU):
+        output_feature = module.net[0].proj.out_features
+        module.net[0].proj.weight.data = torch.cat(
+            [
+                module.net[0].proj.weight.data[: output_feature // 2, :][
+                    lambda_to_keep, :
+                ],
+                module.net[0].proj.weight.data[output_feature // 2 :][
+                    lambda_to_keep, :
+                ],
+            ],
+            dim=0,
+        )
+        module.net[0].proj.out_features = num_lambda * 2
+        if module.net[0].proj.bias is not None:
+            module.net[0].proj.bias.data = torch.cat(
+                [
+                    module.net[0].proj.bias.data[: output_feature // 2][lambda_to_keep],
+                    module.net[0].proj.bias.data[output_feature // 2 :][lambda_to_keep],
+                ]
+            )
+
+        update_act = GEGLU(module.net[0].proj.in_features, num_lambda * 2)
+        update_act.proj = module.net[0].proj
+        module.net[0] = update_act
+
+    # proj weight after activation
+    module.net[2].weight.data = module.net[2].weight.data[:, lambda_to_keep]
+    module.net[2].in_features = num_lambda
+
+    return module
+
+
+# create SparsityLinear module
+class SparsityLinear(torch.nn.Module):
+    def __init__(self, in_features, out_features, lambda_to_keep, num_lambda):
+        super(SparsityLinear, self).__init__()
+        self.linear = torch.nn.Linear(in_features, num_lambda)
+        self.out_features = out_features
+        self.lambda_to_keep = lambda_to_keep
+
+    def forward(self, x):
+        x = self.linear(x)
+        output = torch.zeros(
+            x.size(0), self.out_features, device=x.device, dtype=x.dtype
+        )
+        output[:, self.lambda_to_keep] = x
+        return output
+
+
+def norm_layer_pruning(module, lamb):
+    """
+    Pruning the layer normalization layer for FLUX model
+    """
+    lambda_to_keep = torch.nonzero(lamb).squeeze()
+    if len(lambda_to_keep) == 0:
+        return SkipConnection()
+
+    num_lambda = len(lambda_to_keep)
+
+    # get num_features
+    in_features = module.linear.in_features
+    out_features = module.linear.out_features
+
+    linear = SparsityLinear(in_features, out_features, lambda_to_keep, num_lambda)
+    linear.linear.weight.data = module.linear.weight.data[lambda_to_keep]
+    linear.linear.bias.data = module.linear.bias.data[lambda_to_keep]
+    module.linear = linear
+    return module
+
+
+def get_save_pts(save_pt):
+    if "ff.pt" in save_pt:
+        ff_save_pt = deepcopy(save_pt)  # avoid in-place operation
+        attn_save_pt = save_pt.split(os.sep)
+        attn_save_pt[-1] = attn_save_pt[-1].replace("ff", "attn")
+        attn_save_pt_output = os.sep.join(attn_save_pt)
+        attn_save_pt[-1] = attn_save_pt[-1].replace("attn", "norm")
+        norm_save_pt = os.sep.join(attn_save_pt)
+
+        return {
+            "ff": ff_save_pt,
+            "attn": attn_save_pt_output,
+            "norm": norm_save_pt,
+        }
+    else:
+        attn_save_pt = deepcopy(save_pt)
+        ff_save_pt = save_pt.split(os.sep)
+        ff_save_pt[-1] = ff_save_pt[-1].replace("attn", "ff")
+        ff_save_pt_output = os.sep.join(ff_save_pt)
+        ff_save_pt[-1] = ff_save_pt[-1].replace("ff", "norm")
+        norm_save_pt = os.sep.join(attn_save_pt)
+
+        return {
+            "ff": ff_save_pt_output,
+            "attn": attn_save_pt,
+            "norm": norm_save_pt,
+        }
+
+
+def save_img(pipe, g_cpu, steps, prompt, save_path):
+    image = pipe(prompt, generator=g_cpu, num_inference_steps=steps)
+    image["images"][0].save(save_path)