A FMT block inspried by DiT Block
| 143 | |
| 144 | |
| 145 | class FMTBlock(nn.Module): |
| 146 | """ |
| 147 | A FMT block inspried by DiT Block |
| 148 | """ |
| 149 | def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, **block_kwargs) -> None: |
| 150 | super().__init__() |
| 151 | self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 152 | self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, **block_kwargs) |
| 153 | self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 154 | mlp_hidden_dim = int(hidden_size * mlp_ratio) |
| 155 | approx_gelu = lambda: nn.GELU(approximate="tanh") |
| 156 | self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0) |
| 157 | self.adaLN_modulation = nn.Sequential( |
| 158 | nn.SiLU(), |
| 159 | nn.Linear(hidden_size, 6 * hidden_size, bias=True) |
| 160 | ) |
| 161 | |
| 162 | def framewise_modulate(self, x, shift, scale) -> torch.Tensor: |
| 163 | return x * (1 + scale) + shift |
| 164 | |
| 165 | def forward(self, x, c, mask=None) -> torch.Tensor: |
| 166 | assert mask is not None |
| 167 | shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=-1) |
| 168 | x = x + gate_msa * self.attn(self.framewise_modulate(self.norm1(x), shift_msa, scale_msa), mask = mask) |
| 169 | x = x + gate_mlp * self.mlp(self.framewise_modulate(self.norm2(x), shift_mlp, scale_mlp)) |
| 170 | return x |
| 171 | |
| 172 | class Decoder(nn.Module): |
| 173 | """ |