| 168 | class FinalLayer(Module): |
| 169 | |
| 170 | def __init__(self, |
| 171 | hidden_size, |
| 172 | patch_size, |
| 173 | out_channels, |
| 174 | mapping=Mapping(), |
| 175 | dtype=None): |
| 176 | super().__init__() |
| 177 | self.dtype = dtype |
| 178 | self.norm_final = LayerNorm(hidden_size, |
| 179 | elementwise_affine=False, |
| 180 | eps=1e-6) |
| 181 | self.linear = Linear(hidden_size, |
| 182 | patch_size * patch_size * out_channels, |
| 183 | bias=True, |
| 184 | dtype=dtype) |
| 185 | self.adaLN_modulation = Linear(hidden_size, |
| 186 | 2 * hidden_size, |
| 187 | tp_group=mapping.tp_group, |
| 188 | tp_size=mapping.tp_size, |
| 189 | bias=True, |
| 190 | dtype=dtype) |
| 191 | |
| 192 | def forward(self, x, c): |
| 193 | shift, scale = chunk(self.adaLN_modulation(silu(c)), 2, dim=1) |