Scalable Diffusion Models with Transformers(DIT)代码笔记

完整代码来源：DiT
DiT模型主要是在diffusion中，使用transformer模型替换了UNet模型，使用class来控制图像生成。
根据论文，模型越大，patch size 越小，FID越小。
模型越大，参数越多，patch size越小，参与计算的信息就越多，模型效果越好。

模型使用了Imagenet 训练，有1000个分类，class_labels是0到999的整数，无条件类则必须是1000，在class embedding的时候定义了 nn.Embedding(1000+1,1152)。

在模型初始化的时候，DiTBlock和FinalLayer的参数中weight和bias被置为0，也就是adaLN_zero。
在DiTBlock 中，参数 c = t + y ，也就是 class embedding和 t embedding的和，然后通过线性映射生成shift，scale，和gate。

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

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):
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1)
        x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa))
        x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
        return x


 # Zero-out adaLN modulation layers in DiT blocks:
 for block in self.blocks:
     nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
     nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
     
 # Zero-out output layers:
 nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
 nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
 nn.init.constant_(self.final_layer.linear.weight, 0)
 nn.init.constant_(self.final_layer.linear.bias, 0)