Blame for diffusionGS/models/transformers/dit.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange

try:
    import xformers.ops as xops
except ImportError:
    xops = None

from timm.models.vision_transformer import PatchEmbed, Attention, Mlp

class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))
    def _norm(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
    def forward(self, x):
        output = self._norm(x.float()).type_as(x)
        return output * self.weight.type_as(x)


def modulate(x, shift, scale):
    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)

# https://github.com/karpathy/nanoGPT/blob/eba36e84649f3c6d840a93092cb779a260544d08/model.py#L162-L168
def _init_weights(module):
    if isinstance(module, nn.Linear):
        torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
        if module.bias is not None:
            torch.nn.init.zeros_(module.bias)
    elif isinstance(module, nn.Embedding):
        torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)


class ImageTokenizer(nn.Module):
    """
    Reference: https://github.com/facebookresearch/dino/blob/7c446df5b9f45747937fb0d72314eb9f7b66930a/vision_transformer.py#L134-L214
    """

    def __init__(
        self,
        image_size,
        patch_size,
        d,
        in_channels=3,
        conv_bias=False,
        patch_token_dropout=0.0,
    ):
        super().__init__()
        assert (
            image_size % patch_size == 0
        ), f"Image size {image_size} must be divisible by the patch size {patch_size}."
        num_patches = (image_size // patch_size) * (image_size // patch_size)
        self.image_size = image_size
        self.patch_size = patch_size
        self.num_patches = num_patches
        self.conv = nn.Conv2d(
            in_channels, d, kernel_size=patch_size, stride=patch_size, bias=conv_bias
        )

        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches, d))
        nn.init.trunc_normal_(self.pos_embed, std=0.02)

        self.patch_token_dropout = nn.Dropout(p=patch_token_dropout)

    def forward(self, x):
        # (b, c, h, w) --> (b, l, d)
        # x in [-1, 1]
        x = self.conv(x).flatten(2).transpose(1, 2)
        x = self.patch_token_dropout(x + self.pos_embed)
        return x


class MLP(nn.Module):
    """
    MLP layer
    Reference: https://github.com/facebookresearch/dino/blob/7c446df5b9f45747937fb0d72314eb9f7b66930a/vision_transformer.py#L49-L65
    Ignore: https://github.com/facebookresearch/dino/blob/7c446df5b9f45747937fb0d72314eb9f7b66930a/vision_transformer.py#L62
    """

    def __init__(
        self,
        d,
        mlp_ratio=4,
        mlp_bias=False,
        mlp_dropout=0.0,
        mlp_dim=None,
    ):
        super().__init__()
        if mlp_dim is None:
            mlp_dim = d * mlp_ratio
        self.mlp = nn.Sequential(
            nn.Linear(d, mlp_dim, bias=mlp_bias),
            nn.GELU(),
            nn.Linear(mlp_dim, d, bias=mlp_bias),
            nn.Dropout(mlp_dropout),
        )

    def forward(self, x):
        x = self.mlp(x)
        return x


class SelfAttention(nn.Module):
    """
    Self-attention layer
    Reference: https://github.com/facebookresearch/dino/blob/7c446df5b9f45747937fb0d72314eb9f7b66930a/vision_transformer.py#L68-L92
    """

    def __init__(
        self,
        d,
        d_head,
        attn_qkv_bias=False,
        attn_dropout=0.0,
        attn_fc_bias=False,
        attn_fc_dropout=0.0,
        use_flashatt_v2=True,
    ):
        super().__init__()
        assert (
            d % d_head == 0
        ), f"Token dimension {d} should be divisible by head dimension {d_head}"
        self.d = d
        self.d_head = d_head
        self.attn_dropout = attn_dropout

        self.to_qkv = nn.Linear(d, 3 * d, bias=attn_qkv_bias)
        self.fc = nn.Linear(d, d, bias=attn_fc_bias)
        self.attn_fc_dropout = nn.Dropout(attn_fc_dropout)

        self.use_flashatt_v2 = use_flashatt_v2 and xops is not None

    def forward(self, x, subset_attention_size=None):
        """
        x: (b, l, d)
        """
        # token split, multi-head attention, token cat
        q, k, v = self.to_qkv(x).split(self.d, dim=2)

        if self.use_flashatt_v2:
            # Use the flash attention support from the xformers library
            # The memory_efficient_attention takes the input as (batch, seq_len, heads, dim)
            q, k, v = map(
                lambda t: rearrange(t, "b l (nh dh) -> b l nh dh", dh=self.d_head),
                (q, k, v),
            )

            if subset_attention_size is not None and subset_attention_size < q.shape[1]:
                x_subset = xops.memory_efficient_attention(
                    q[:, :subset_attention_size, :, :].contiguous(),
                    k[:, :subset_attention_size, :, :].contiguous(),
                    v[:, :subset_attention_size, :, :].contiguous(),
                    attn_bias=None,
                    op=(xops.fmha.flash.FwOp, xops.fmha.flash.BwOp),
                )
                x_rest = xops.memory_efficient_attention(
                    q[:, subset_attention_size:, :, :].contiguous(),
                    k,
                    v,
                    attn_bias=None,
                    op=(xops.fmha.flash.FwOp, xops.fmha.flash.BwOp),
                )
                # print(x_subset.shape, x_rest.shape)
                x = torch.cat([x_subset, x_rest], dim=1)
            else:
                # The memory_efficient_attention takes the input as (batch, seq_len, heads, dim)
                x = xops.memory_efficient_attention(
                    q,
                    k,
                    v,
                    attn_bias=None,
                    op=(xops.fmha.flash.FwOp, xops.fmha.flash.BwOp),
                )
            x = rearrange(x, "b l nh dh -> b l (nh dh)")
        else:
            # Use the flash attention support from the pytorch library
            q, k, v = (
                rearrange(q, "b l (nh dh) -> b nh l dh", dh=self.d_head),
                rearrange(k, "b l (nh dh) -> b nh l dh", dh=self.d_head),
                rearrange(v, "b l (nh dh) -> b nh l dh", dh=self.d_head),
            )
            # https://discuss.pytorch.org/t/flash-attention/174955/14
            dropout_p = self.attn_dropout if self.training else 0.0
            if subset_attention_size is not None and subset_attention_size < q.shape[2]:
                x = F.scaled_dot_product_attention(
                    q[:, :, :subset_attention_size, :].contiguous(),
                    k[:, :, :subset_attention_size, :].contiguous(),
                    v[:, :, :subset_attention_size, :].contiguous(),
                    dropout_p=dropout_p,
                )
                x_rest = F.scaled_dot_product_attention(
                    q[:, :, subset_attention_size:, :].contiguous(),
                    k,
                    v,
                    dropout_p=dropout_p,
                )
                x = torch.cat([x, x_rest], dim=2)
            else:
                x = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p)
            x = rearrange(x, "b nh l dh -> b l (nh dh)")

        x = self.attn_fc_dropout(self.fc(x))
        return x


class TransformerBlock(nn.Module):
    """
    Transformer block
    Reference: https://github.com/facebookresearch/dino/blob/7c446df5b9f45747937fb0d72314eb9f7b66930a/vision_transformer.py#L95-L113
    Note: move drop_path to SelfAttention and MLP
    """

    def __init__(
        self,
        d,
        d_head,
        ln_bias=False,
        attn_qkv_bias=False,
        attn_dropout=0.0,
        attn_fc_bias=False,
        attn_fc_dropout=0.0,
        mlp_ratio=4,
        mlp_bias=False,
        mlp_dropout=0.0,
    ):
        super().__init__()
        self.norm1 = nn.LayerNorm(d, bias=ln_bias)
        self.attn = SelfAttention(
            d, d_head, attn_qkv_bias, attn_dropout, attn_fc_bias, attn_fc_dropout
        )
        self.norm2 = nn.LayerNorm(d, bias=ln_bias)
        self.mlp = MLP(d, mlp_ratio, mlp_bias, mlp_dropout)

    def forward(self, x, subset_attention_size=None):
        x = x + self.attn(self.norm1(x), subset_attention_size=subset_attention_size)
        x = x + self.mlp(self.norm2(x))
        return x



class DiTBlock(nn.Module):
    """
    A DiT block with adaptive layer norm zero (adaLN-Zero) conditioning.
    """

    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, **block_kwargs):
        super().__init__()
        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.attn = Attention(
            hidden_size, num_heads=num_heads, qkv_bias=True, **block_kwargs
        )
        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        mlp_hidden_dim = int(hidden_size * mlp_ratio)
        approx_gelu = lambda: nn.GELU(approximate="tanh")
        self.mlp = Mlp(
            in_features=hidden_size,
            hidden_features=mlp_hidden_dim,
            act_layer=approx_gelu,
            drop=0,
        )
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True)
        )

    def forward(self, x, c):
        '''
            x 是 latent feature, 已经经过了 patch embedding
            c 是 conditioning, 主要有两部分 - timestep embedding 和 label embedding
            此处就只有一个 timestep embedding
        '''
        # 这算作 conditioning 的 MLP
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
            self.adaLN_modulation(c).chunk(6, dim=1)
        )   # 将 c 通过 adaLN_modulation 网络, 分成 6 个部分
        # 第一个模块 - MSA
        x = x + gate_msa.unsqueeze(1) * self.attn(
            modulate(self.norm1(x), shift_msa, scale_msa)
        )

        # 第二个模块 - Pointwise Feedforward
        x = x + gate_mlp.unsqueeze(1) * self.mlp(
            modulate(self.norm2(x), shift_mlp, scale_mlp)
        )
        return x



class CrossAttention(nn.Module):
    """
    Self-attention layer
    Reference: https://github.com/facebookresearch/dino/blob/7c446df5b9f45747937fb0d72314eb9f7b66930a/vision_transformer.py#L68-L92
    """

    def __init__(
        self,
        input_dim,
        d_head=64,
        attn_qkv_bias=False,
        attn_dropout=0.0,
        attn_fc_bias=False,
        attn_fc_dropout=0.0,
        use_flashatt_v2=True,
        num_heads=None,
        ctx_dim=None,
        causal=False,
    ):
        super().__init__()
        self.input_dim = input_dim
        self.d_head = d_head
        self.num_heads = num_heads if num_heads is not None else input_dim // d_head
        self.ctx_dim = ctx_dim if ctx_dim is not None else input_dim
        self.att_dim = self.num_heads * self.d_head

        self.to_q = nn.Linear(self.input_dim, self.att_dim, bias=attn_qkv_bias)
        self.to_k = nn.Linear(self.ctx_dim, self.att_dim, bias=attn_qkv_bias)
        self.to_v = nn.Linear(self.ctx_dim, self.att_dim, bias=attn_qkv_bias)
        self.fc = nn.Linear(self.att_dim, self.input_dim, bias=attn_fc_bias)
        self.attn_fc_dropout = nn.Dropout(attn_fc_dropout)

        self.attn_dropout = attn_dropout
        assert self.attn_dropout == 0.0
        self.use_flashatt_v2 = use_flashatt_v2 and xops is not None
        self.causal = causal

    def forward(self, x, y=None):
        """
        x: (b, l, d)
        y: (b, l', d)
        x cross-att y
        """
        if y is None:
            y = x

        q = self.to_q(x)
        k = self.to_k(y)
        v = self.to_v(y)

        if self.use_flashatt_v2:
            q, k, v = map(
                lambda t: rearrange(t, "b n (h d) -> b n h d", h=self.num_heads),
                (q, k, v),
            )

            # Use the flash attention support from the xformers library
            if self.causal:
                attention_bias = xops.LowerTriangularMask()
            else:
                attention_bias = None

            # The memory_efficient_attention takes the input as (batch, seq_len, heads, dim)
            x = xops.memory_efficient_attention(
                q,
                k,
                v,
                attn_bias=attention_bias,
                op=(xops.fmha.flash.FwOp, xops.fmha.flash.BwOp),
            )

            x = rearrange(x, "b n h d -> b n (h d)")
        else:
            q, k, v = map(
                lambda t: rearrange(t, "b n (h d) -> b h n d", h=self.num_heads),
                (q, k, v),
            )

            # https://discuss.pytorch.org/t/flash-attention/174955/14
            dropout_p = self.attn_dropout if self.training else 0.0
            x = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p)

            x = rearrange(x, "b nh l dh -> b l (nh dh)")

        x = self.attn_fc_dropout(self.fc(x))
        return x

    def extra_repr(self) -> str:
        return (
            f"use_flashatt_v2={self.use_flashatt_v2}, "
            f"num_heads={self.num_heads}, "
            f"input_dim={self.input_dim}, "
            f"ctx_dim={self.ctx_dim}, "
            f"att_dim={self.att_dim}, "
        )


class CrossTransformerBlock(nn.Module):
    """
    Cross-att Transformer block
    Reference: https://github.com/facebookresearch/dino/blob/7c446df5b9f45747937fb0d72314eb9f7b66930a/vision_transformer.py#L95-L113
    Note: move drop_path to SelfAttention and MLP
    """

    def __init__(
        self,
        d,
        d_head,
        ln_bias=False,
        attn_qkv_bias=False,
        attn_dropout=0.0,
        attn_fc_bias=False,
        attn_fc_dropout=0.0,
        mlp_ratio=4,
        mlp_bias=False,
        mlp_dropout=0.0,
    ):
        super().__init__()
        self.norm1 = nn.LayerNorm(d, bias=ln_bias)

        self.attn = CrossAttention(
            d, d_head, attn_qkv_bias, attn_dropout, attn_fc_bias, attn_fc_dropout
        )
        self.norm2 = nn.LayerNorm(d, bias=ln_bias)
        self.mlp = MLP(d, mlp_ratio, mlp_bias, mlp_dropout)

    def forward(self, x, y):
        x = x + self.attn(self.norm1(x), self.norm1(y))
        x = x + self.mlp(self.norm2(x))
        return x


class FixedLengthTransformerLayer(nn.Module):
    """ """

    def __init__(
        self,
        dim,
        context_dim=None,
        fixed_length=None,
        num_heads=8,
        head_dim=64,
        use_ln_context=True,
        mlp_dim=None,
    ):
        """
        :param dim: The input dim of x
        :param context_dim: The input dim of context
        :param fixed_length: The length of attention tokens
        :param num_heads: The number of attention heads
        :param head_dim: The dim of each attention head
        """
        super().__init__()
        self.has_cross_att = context_dim is not None
        self.dim = dim
        self.fixed_length = fixed_length

        self.ln_self = nn.LayerNorm(dim)
        self.self_attn = CrossAttention(
            input_dim=dim,
            d_head=head_dim,
            num_heads=num_heads,
        )

        if self.has_cross_att:
            self.ln_cross = nn.LayerNorm(dim)
            if use_ln_context:
                self.ln_context = nn.LayerNorm(context_dim)
            else:
                self.ln_context = nn.Identity()
            self.cross_attn = CrossAttention(
                input_dim=dim,
                ctx_dim=context_dim,
                d_head=head_dim,
                num_heads=num_heads,
            )

        self.ln_fc = nn.LayerNorm(dim)
        self.fc = MLP(
            d=dim,
            mlp_dim=mlp_dim,
        )

    def init_weight(self, total_layers):
        # Find all the linear layers that will contribute to the residual.
        linear_layer_list = [
            self.self_attn.fc,
            self.fc.mlp[2],
        ]
        if self.has_cross_att:
            linear_layer_list.append(self.cross_attn.fc)

        # divided the output by the total number of layers
        for linear_layer in linear_layer_list:
            linear_layer.weight.data /= total_layers
            if linear_layer.bias is not None:
                linear_layer.bias.data = 0.0

    def forward(self, x, context=None):
        """
        :param x: b, h_windows, w_windows, window_size, window_size, c;
                    can be in arbitrary shape but the layout is the same.
        :param context: b, context_length, context_dim
        :return:
        """
        batch_size, orig_length, orig_dim = x.shape
        context = context or x

        assert orig_dim % self.dim == 0, f"orig_dim: {orig_dim}, dim: {self.dim}"
        if self.fixed_length is not None:
            assert (orig_length * orig_dim) % (self.fixed_length * self.dim) == 0, (
                f"orig_length: {orig_length}, token_length: {self.fixed_length}"
                f"orig_dim: {orig_dim}, dim: {self.dim}."
                f"The product of orig_length * orig_dim must be divisible by token_length * dim."
                f"O.w., it will break the batches"
            )

        # Fixed-length self attention [-1, highres_tokens, highres_dim], e.g., [-1, 8192, 16]
        x = x.reshape(-1, self.fixed_length or orig_length, self.dim)
        x = x + self.self_attn(self.ln_self(x))

        # Cross attention [batch_size, -1, highres_dim], e.g., [8, -1, 16]
        if self.has_cross_att:
            x = x.reshape(batch_size, -1, self.dim)
            x = x + self.cross_attn(self.ln_cross(x), self.ln_context(context))

        # MLP layers [?, ?, highres_dim], e.g., [?, ?, 16]
        x = x + self.fc(self.ln_fc(x))

        # Reshape back to [batch_size, lowres_tokens, lowres_dim], e.g., [8, 1024, 1024]
        x = x.reshape(batch_size, orig_length, orig_dim)

        return x

    def extra_repr(self) -> str:
        return f"token_length={self.fixed_length}, dim={self.dim}"



class QK_Norm_SelfAttention(nn.Module):
    """
    Self-attention layer
    Reference: https://github.com/facebookresearch/dino/blob/7c446df5b9f45747937fb0d72314eb9f7b66930a/vision_transformer.py#L68-L92
    """
    def __init__(
        self,
        d,
        d_head,
        attn_qkv_bias=False,
        attn_fc_bias=True,
        attn_dropout=0.0,
        attn_fc_dropout=0.0,
        use_qk_norm=True,
    ):
        super().__init__()
        assert d % d_head == 0, f"Token dimension {d} should be divisible by head dimension {d_head}"
        self.d = d
        self.d_head = d_head
        self.attn_dropout = attn_dropout
        self.to_qkv = nn.Linear(d, 3 * d, bias=attn_qkv_bias)
        self.fc = nn.Linear(d, d, bias=attn_fc_bias)
        self.attn_fc_dropout = nn.Dropout(attn_fc_dropout)
        self.use_qk_norm = use_qk_norm
        if self.use_qk_norm:
            self.q_norm = RMSNorm(d_head)
            self.k_norm = RMSNorm(d_head)
    def forward(self, x, attn_bias=None):
        """
        x: (b, l, d)
        attn_bias: xformers BlockDiagonalMask
        """
        # token split, multi-head attention, token cat
        q, k, v = self.to_qkv(x).split(self.d, dim=2)
        # Use the flash attention support from the xformers library
        # The memory_efficient_attention takes the input as (batch, seq_len, heads, dim)
        q, k, v = (rearrange(t, "b l (nh dh) -> b l nh dh", dh=self.d_head) for t in (q, k, v))
        if self.use_qk_norm:
            q = self.q_norm(q)
            k = self.k_norm(k)
        # The memory_efficient_attention takes the input as (batch, seq_len, heads, dim)
        x = xops.memory_efficient_attention(
            q,
            k,
            v,
            attn_bias=attn_bias,
            p=self.attn_dropout if self.training else 0.0,
            op=(xops.fmha.flash.FwOp, xops.fmha.flash.BwOp),
        )
        x = rearrange(x, "b l nh dh -> b l (nh dh)")
        x = self.attn_fc_dropout(self.fc(x))
        return x




# 修改，加入 timestep condition 来加以控制
class DiTBlock_QK_Norm(nn.Module):
    """
    Transformer block
    Reference: https://github.com/facebookresearch/dino/blob/7c446df5b9f45747937fb0d72314eb9f7b66930a/vision_transformer.py#L95-L113
    Note: move drop_path to SelfAttention and MLP
    """
    def __init__(
        self,
        d,
        d_head,
        ln_bias=True,
        attn_qkv_bias=False,
        attn_dropout=0.0,
        attn_fc_bias=True,
        attn_fc_dropout=0.0,
        mlp_ratio=4,
        mlp_bias=True,
        mlp_dropout=0.0,
        use_qk_norm=True,
    ):
        super().__init__()
        self.norm1 = nn.LayerNorm(d, bias=ln_bias)
        self.attn = QK_Norm_SelfAttention(
            d, d_head, attn_qkv_bias, attn_fc_bias, attn_dropout, attn_fc_dropout, use_qk_norm,
        )
        self.norm2 = nn.LayerNorm(d, bias=ln_bias)
        self.mlp = MLP(d, mlp_ratio, mlp_bias, mlp_dropout)
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(), nn.Linear(d, 6 * d, bias=True)
        )
    def forward(self, x, c):
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
            self.adaLN_modulation(c).chunk(6, dim=1)
        )   # 将 c 通过 adaLN_modulation 网络, 分成 6 个部分

        x = x + gate_msa.unsqueeze(1) * self.attn(
            modulate(self.norm1(x), shift_msa, scale_msa)
        )

        # 第二个模块 - Pointwise Feedforward
        x = x + gate_mlp.unsqueeze(1) * self.mlp(
            modulate(self.norm2(x), shift_mlp, scale_mlp)
        )
        return x