| 216 | |
| 217 | |
| 218 | class FinalLayer(nn.Module): |
| 219 | def __init__(self, model): |
| 220 | super().__init__() |
| 221 | self.norm_out = model.norm_out |
| 222 | self.proj_out = model.proj_out |
| 223 | self.model = [model] |
| 224 | |
| 225 | def __getattr__(self, name): |
| 226 | return getattr(self.model[0], name) |
| 227 | |
| 228 | @torch.autocast('cuda', dtype=AUTOCAST_DTYPE) |
| 229 | def forward(self, inputs): |
| 230 | hidden_states, temb, latent_size, *_ = inputs |
| 231 | height = latent_size[0].item() |
| 232 | width = latent_size[1].item() |
| 233 | |
| 234 | hidden_states = self.norm_out(hidden_states, temb) |
| 235 | hidden_states = self.proj_out(hidden_states) |
| 236 | |
| 237 | # unpatchify |
| 238 | patch_size = self.config.patch_size |
| 239 | height = height // patch_size |
| 240 | width = width // patch_size |
| 241 | |
| 242 | hidden_states = hidden_states.reshape( |
| 243 | shape=(hidden_states.shape[0], height, width, patch_size, patch_size, self.out_channels) |
| 244 | ) |
| 245 | hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) |
| 246 | output = hidden_states.reshape( |
| 247 | shape=(hidden_states.shape[0], self.out_channels, height * patch_size, width * patch_size) |
| 248 | ) |
| 249 | return output |