add ecodiff demo

app.py

@@ -1,154 +1,224 @@
 import gradio as gr
-import numpy as np
-import random
+from dataclasses import dataclass
 
-import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
 import torch
+from tqdm import tqdm
+
+from src.utils import (
+    create_pipeline,
+    calculate_mask_sparsity,
+    ffn_linear_layer_pruning,
+    linear_layer_pruning,
+)
+from diffusers import StableDiffusionXLPipeline
+
+
+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"
+    device: str = "cuda:0"
+    seed: list = None
+    prompt: str = None
+    mix_precision: str = "bf16"
+    num_intervention_steps: int = 50
+    model: str = "sdxl"
+    binary: bool = False
+    masking: str = "binary"
+    scope: str = "global"
+    ratio: list = None
+    width: int = None
+    height: int = None
+    epsilon: float = 0.0
+    lambda_threshold: float = 0.001
+
+    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(args.device)
+
+    device = args.device
+    torch_dtype = torch.bfloat16 if args.mix_precision == "bf16" else torch.float32
+    mask_pipe, hookers = create_pipeline(
+        pipe,
+        args.model,
+        device,
+        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,
+    )
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
-
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
-
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
-
-
-@spaces.GPU #[uncomment to use ZeroGPU]
-def infer(
-    prompt,
-    negative_prompt,
-    seed,
-    randomize_seed,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-    progress=gr.Progress(track_tqdm=True),
-):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    generator = torch.Generator().manual_seed(seed)
-
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
-
-    return image, seed
-
-
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
-]
-
-css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 640px;
-}
-"""
-
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # Text-to-Image Gradio Template")
+    # 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}"
+        )
 
-        with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
+    # Prune the model
+    pruned_pipe = prune_model(mask_pipe, hookers)
 
-            run_button = gr.Button("Run", scale=0, variant="primary")
+    # reload the original model
+    pipe = StableDiffusionXLPipeline.from_pretrained(
+        "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16
+    ).to(args.device)
 
-        result = gr.Image(label="Result", show_label=False)
+    # 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
+
+
+def generate_images(prompt, seed, steps, pipe, pruned_pipe):
+    # Run the model and return images directly
+    g_cpu = torch.Generator("cuda:0").manual_seed(seed)
+    original_image = pipe(prompt=prompt, generator=g_cpu, num_inference_steps=steps).images[0]
+    g_cpu = torch.Generator("cuda:0").manual_seed(seed)
+    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):
+    args = GradioArgs(prompt=prompt, seed=[seed], num_intervention_steps=steps)
+    pipe, pruned_pipe = binary_mask_eval(args)
+    return pipe, pruned_pipe, [("Model Initialized", "green")]
 
-        with gr.Accordion("Advanced Settings", open=False):
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
-            )
 
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
+def on_generate_click(prompt, seed, steps, pipe, pruned_pipe):
+    original_image, ecodiff_image = generate_images(prompt, seed, steps, pipe, pruned_pipe)
+    return original_image, ecodiff_image
+
+
+def create_demo():
+    with gr.Blocks() as demo:
+        gr.Markdown("# Text-to-Image Generation with EcoDiff Pruned Model")
+        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.**
+                """
             )
+        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)
+            prune_btn = gr.Button("Initialize Original and Pruned Models", variant="primary", scale=1)
+            status_label = gr.HighlightedText(label="Model Status", value=[("Model Not Initialized", "red")], scale=1)
+        with gr.Row():
+            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)
+            generate_btn = gr.Button("Generate Images")
+        gr.Examples(
+            examples=[
+                "A clock tower floating in a sea of clouds",
+                "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
+                "An astronaut riding a green horse",
+                "A delicious ceviche cheesecake slice",
+            ],
+            inputs=[prompt],
+        )
+        with gr.Row():
+            original_output = gr.Image(label="Original Output")
+            ecodiff_output = gr.Image(label="EcoDiff Output")
+
+        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],
+            outputs=[original_output, ecodiff_output],
+        )
+        prune_btn.click(
+            fn=on_prune_click,
+            inputs=[prompt, seed, steps],
+            outputs=[pipe_state, pruned_pipe_state, status_label],
+        )
+        generate_btn.click(
+            fn=on_generate_click,
+            inputs=[prompt, seed, steps, pipe_state, pruned_pipe_state],
+            outputs=[original_output, ecodiff_output],
+        )
+
+    return demo
 
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-            with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
-
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
-
-            with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=0.0,  # Replace with defaults that work for your model
-                )
-
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=2,  # Replace with defaults that work for your model
-                )
-
-        gr.Examples(examples=examples, inputs=[prompt])
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
-        inputs=[
-            prompt,
-            negative_prompt,
-            seed,
-            randomize_seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-        ],
-        outputs=[result, seed],
-    )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo = create_demo()
+    demo.launch(share=True)

requirements.txt

@@ -1,6 +1,4 @@
-accelerate
-diffusers
-invisible_watermark
-torch
-transformers
-xformers
+diffusers==0.31.0
+torch==2.4.1
+transformers==4.45.2
+accelerate==0.33.0

src/cross_attn_hook.py

@@ -0,0 +1,425 @@
+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
+
+
+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 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,
+        )
+        self.attention_processor = AttnProcessor2_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
+
+            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
+
+        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

src/ffn_hooker.py

@@ -0,0 +1,193 @@
+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

src/utils.py

@@ -0,0 +1,245 @@
+import os
+from copy import deepcopy
+from typing import Optional
+
+import torch
+from diffusers.models.activations import GEGLU, GELU
+from src.cross_attn_hook import CrossAttentionExtractionHook
+from src.ffn_hooker import FeedForwardHooker
+
+
+# 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)
+            norm_hooker = None
+
+            g_cpu = torch.Generator(torch.device(device)).manual_seed(1)
+            _ = pipe("abc", generator=g_cpu, 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
+
+
+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)