| 310 | The final layer of LightningDiT. |
| 311 | """ |
| 312 | def __init__(self, hidden_size, patch_size, out_channels, use_rmsnorm=False): |
| 313 | super().__init__() |
| 314 | if not use_rmsnorm: |
| 315 | self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 316 | else: |
| 317 | self.norm_final = RMSNorm(hidden_size) |
| 318 | self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True) |
| 319 | self.adaLN_modulation = nn.Sequential( |
| 320 | nn.SiLU(), |
| 321 | nn.Linear(hidden_size, 2 * hidden_size, bias=True) |
| 322 | ) |
| 323 | @torch.compile |
| 324 | def forward(self, x, c): |
| 325 | shift, scale = self.adaLN_modulation(c).chunk(2, dim=1) |