Understand Attention in SD

class AttnProcessor:  # 定义一个名为 AttnProcessor 的类
    r"""
    用于执行注意力相关计算的默认处理器。
    """

    def __call__(  # 定义该类的可调用方法，使其像函数一样被调用
        self,
        attn: Attention,  # 传入的 Attention 模块实例
        hidden_states,  # 输入的隐藏状态张量 (通常是图像特征)
        encoder_hidden_states=None,  # 用于交叉注意力的编码器隐藏状态，如果为 None 则是自注意力
        attention_mask=None,  # 注意力掩码，用于忽略某些位置的计算
        temb=None,  # 时间步嵌入 (Time Embedding)，通常用于条件扩散模型
    ):
        residual = hidden_states  # 保存原始的 hidden_states，用于之后的残差连接

        if attn.spatial_norm is not None:  # 检查是否存在空间归一化层 (spatial_norm)
            hidden_states = attn.spatial_norm(
                hidden_states, temb
            )  # 如果存在，则应用空间归一化

        input_ndim = hidden_states.ndim  # 获取输入 hidden_states 张量的维度数

        if (
            input_ndim == 4
        ):  # 如果输入是4维张量 (通常是 [batch, channel, height, width])
            batch_size, channel, height, width = (
                hidden_states.shape
            )  # 获取张量的各个维度大小
            hidden_states = hidden_states.view(  # 将4D张量重塑为3D
                batch_size, channel, height * width
            ).transpose(
                1, 2
            )  # 并交换第1和第2个维度，变为 [batch, sequence_length, channel]

        # 根据是自注意力还是交叉注意力，获取批次大小和序列长度
        batch_size, sequence_length, _ = (
            hidden_states.shape  # 如果是自注意力，从 hidden_states 获取形状
            if encoder_hidden_states is None  # 判断条件：encoder_hidden_states 是否为空
            else encoder_hidden_states.shape  # 如果是交叉注意力，从 encoder_hidden_states 获取形状
        )
        # 准备注意力掩码，确保其形状和设备与计算兼容
        attention_mask = attn.prepare_attention_mask(
            attention_mask, sequence_length, batch_size
        )

        if attn.group_norm is not None:  # 检查是否存在组归一化层 (group_norm)
            # 应用组归一化。需要先将通道维度换回来，归一化后再换回去
            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(
                1, 2
            )

        query = attn.to_q(
            hidden_states
        )  # 通过一个线性层将 hidden_states 映射为 Query (Q)

        if encoder_hidden_states is None:  # 如果是自注意力 (self-attention)
            encoder_hidden_states = hidden_states  # Key 和 Value 的来源与 Query 相同
        elif attn.norm_cross:  # 如果是交叉注意力 (cross-attention) 且需要归一化
            encoder_hidden_states = (
                attn.norm_encoder_hidden_states(  # 对 encoder_hidden_states 进行归一化
                    encoder_hidden_states
                )
            )

        key = attn.to_k(
            encoder_hidden_states
        )  # 通过线性层将 encoder_hidden_states 映射为 Key (K)
        value = attn.to_v(
            encoder_hidden_states
        )  # 通过线性层将 encoder_hidden_states 映射为 Value (V)

        query = attn.head_to_batch_dim(
            query
        )  # 为多头注意力机制重塑 Q，将 "头" 的维度合并到 "批次" 维度
        key = attn.head_to_batch_dim(key)  # 为多头注意力机制重塑 K
        value = attn.head_to_batch_dim(value)  # 为多头注意力机制重塑 V

        # 计算注意力分数 (通常是 Q 和 K 的点积，然后应用 softmax)
        attention_probs = attn.get_attention_scores(query, key, attention_mask)
        # 将注意力分数与 V 相乘，得到加权和，即注意力的输出
        hidden_states = torch.bmm(attention_probs, value)
        # 将 "头" 和 "批次" 维度分离，恢复原始的张量结构
        hidden_states = attn.batch_to_head_dim(hidden_states)

        # linear proj (线性投影)
        hidden_states = attn.to_out[0](hidden_states)  # 通过输出线性层进行投影
        # dropout
        hidden_states = attn.to_out[1](hidden_states)  # 应用 Dropout 防止过拟合

        if input_ndim == 4:  # 如果原始输入是4维的
            # 将张量恢复为原始的4D图像格式 [batch, channel, height, width]
            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  # 返回最终处理后的隐藏状态

# 计算注意力分数 (通常是 Q 和 K 的点积，然后应用 softmax)
attention_probs = attn.get_attention_scores(query, key, attention_mask)

class CrossAttnProcessor:  # 定义一个名为 AttnProcessor 的类
    r"""
    用于执行注意力相关计算的默认处理器。
    """

    def __init__(self, attention_store, place_in_unet):
        self.attnstore = attention_store
        self.place_in_unet = place_in_unet
    ...
    def __call__(  # 定义该类的可调用方法，使其像函数一样被调用
        self,
        attn: Attention,  # 传入的 Attention 模块实例
        hidden_states,  # 输入的隐藏状态张量 (通常是图像特征)
        encoder_hidden_states=None,  # 用于交叉注意力的编码器隐藏状态，如果为 None 则是自注意力
        attention_mask=None,  # 注意力掩码，用于忽略某些位置的计算
        temb=None,  # 时间步嵌入 (Time Embedding)，通常用于条件扩散模型
    ):
    ...

        # 计算注意力分数 (通常是 Q 和 K 的点积，然后应用 softmax)
        attention_probs = attn.get_attention_scores(query, key, attention_mask)
        
        self.attnstore(attention_probs, is_cross=True, place_in_unet=self.place_in_unet)
    ...

class AttentionStore:
    """一个简单的注意力图存储类"""

    @staticmethod
    def get_empty_store():
        return {"down_cross": [], "mid_cross": [], "up_cross": []}

    def __call__(self, attn, is_cross: bool, place_in_unet: str, **kwargs):
        # attn.shape: [batch_size * num_heads, seq_len_q, seq_len_k]
        if is_cross:
            key = f"{place_in_unet}_cross"
            # 为了演示，我们只存储注意力图
            self.step_store[key].append(attn.detach().cpu())

    def __init__(self):
        self.step_store = self.get_empty_store()
        self.attention_store = {}

attn_procs = {}
for name in unet.attn_processors.keys():
    # 判断模块在UNet中的位置
    if name.startswith("mid_block"):
        place_in_unet = "mid"
    elif name.startswith("up_blocks"):
        place_in_unet = "up"
    elif name.startswith("down_blocks"):
        place_in_unet = "down"
    else:
        continue

    # 只替换交叉注意力模块 (attn2)
    if "attn2" in name:
        print(f"  - Replacing processor for: {name}")
        attn_procs[name] = CrossAttnProcessor(
            attention_store=attention_store, place_in_unet=place_in_unet
        )
    else:
        # 其他模块（如自注意力attn1）保持默认处理器
        attn_procs[name] = AttnProcessor()

# c. 执行替换
print("\nSetting new attention processors on UNet...")
unet.set_attn_processor(attn_procs)