The final layer of DiT.
| 311 | |
| 312 | |
| 313 | class FinalLayer(nn.Module): |
| 314 | """ |
| 315 | The final layer of DiT. |
| 316 | """ |
| 317 | |
| 318 | def __init__(self, hidden_size, patch_size, out_channels, cond=False): |
| 319 | super().__init__() |
| 320 | self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 321 | self.linear = nn.Linear( |
| 322 | hidden_size, patch_size * patch_size * out_channels, bias=True |
| 323 | ) |
| 324 | if cond: |
| 325 | self.adaLN_modulation = nn.Sequential( |
| 326 | nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True) |
| 327 | ) |
| 328 | |
| 329 | def forward(self, x, c=None): |
| 330 | if c is not None: |
| 331 | shift, scale = self.adaLN_modulation(c).chunk(2, dim=1) |
| 332 | x = modulate(self.norm_final(x), shift, scale) |
| 333 | x = self.linear(x) |
| 334 | else: |
| 335 | x = self.norm_final(x) |
| 336 | x = self.linear(x) |
| 337 | return x |
| 338 | |
| 339 | |
| 340 | class Block(nn.Module): |