The final layer of DiT.
| 233 | |
| 234 | |
| 235 | class FinalLayer(nn.Module): |
| 236 | """ |
| 237 | The final layer of DiT. |
| 238 | """ |
| 239 | def __init__(self, hidden_size, patch_size, out_channels): |
| 240 | super().__init__() |
| 241 | self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 242 | self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True) |
| 243 | self.adaLN_modulation = nn.Sequential( |
| 244 | nn.SiLU(), |
| 245 | nn.Linear(hidden_size, 2 * hidden_size, bias=True) |
| 246 | ) |
| 247 | |
| 248 | def forward(self, x, c): |
| 249 | shift, scale = self.adaLN_modulation(c).chunk(2, dim=1) |
| 250 | x = modulate(self.norm_final(x), shift, scale) |
| 251 | x = self.linear(x) |
| 252 | return x |
| 253 | |
| 254 | |
| 255 | def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0, interpolation_scale=1.0, base_size=1): |