| 490 | |
| 491 | |
| 492 | class FinalLayer(nn.Module): |
| 493 | def __init__(self, model): |
| 494 | super().__init__() |
| 495 | self.final_layer = model.final_layer |
| 496 | self.model = [model] |
| 497 | |
| 498 | def __getattr__(self, name): |
| 499 | return getattr(self.model[0], name) |
| 500 | |
| 501 | @torch.autocast('cuda', dtype=AUTOCAST_DTYPE) |
| 502 | def forward(self, inputs): |
| 503 | img, txt, pe, mod_vectors, txt_img_mask = inputs |
| 504 | img = img[:, txt.shape[1] :, ...] |
| 505 | final_mod_spec = modulation_distribute_dict["final_layer.adaLN_modulation.1"] |
| 506 | final_mod = [mod_vectors[:, s, :] for s in final_mod_spec] |
| 507 | img = self.final_layer( |
| 508 | img, distill_vec=final_mod |
| 509 | ) # (N, T, patch_size ** 2 * out_channels) |
| 510 | return img |