| 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 |