The final layer adopted from DiT.
| 129 | |
| 130 | |
| 131 | class FinalLayer(nn.Module): |
| 132 | """ |
| 133 | The final layer adopted from DiT. |
| 134 | """ |
| 135 | def __init__(self, model_channels, out_channels): |
| 136 | super().__init__() |
| 137 | self.norm_final = nn.LayerNorm(model_channels, elementwise_affine=False, eps=1e-6) |
| 138 | self.linear = nn.Linear(model_channels, out_channels, bias=True) |
| 139 | self.adaLN_modulation = nn.Sequential( |
| 140 | nn.SiLU(), |
| 141 | nn.Linear(model_channels, 2 * model_channels, bias=True) |
| 142 | ) |
| 143 | |
| 144 | def forward(self, x, c): |
| 145 | shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1) |
| 146 | x = modulate(self.norm_final(x), shift, scale) |
| 147 | x = self.linear(x) |
| 148 | return x |
| 149 | |
| 150 | |
| 151 | class SimpleMLPAdaLN(nn.Module): |