| 488 | The final layer of DiT. |
| 489 | """ |
| 490 | def __init__(self, hidden_size, patch_size, out_channels): |
| 491 | super().__init__() |
| 492 | self.norm_final = nn.LayerNorm( |
| 493 | hidden_size, elementwise_affine=False, eps=1e-6, |
| 494 | ) |
| 495 | self.linear = ColumnParallelLinear( |
| 496 | hidden_size, patch_size * patch_size * out_channels, bias=True, |
| 497 | init_method=nn.init.zeros_, gather_output=True, |
| 498 | ) |
| 499 | self.adaLN_modulation = nn.Sequential( |
| 500 | nn.SiLU(), |
| 501 | ColumnParallelLinear( |
| 502 | min(hidden_size, 1024), 2 * hidden_size, bias=True, |
| 503 | init_method=nn.init.zeros_, gather_output=True, |
| 504 | ), |
| 505 | ) |
| 506 | |
| 507 | def forward(self, x, c): |
| 508 | shift, scale = self.adaLN_modulation(c).chunk(2, dim=1) |