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