| 614 | |
| 615 | |
| 616 | class FinalLayer(nn.Module): |
| 617 | def __init__(self, model): |
| 618 | super().__init__() |
| 619 | self.final_layer = model.final_layer |
| 620 | self.model = [model] |
| 621 | |
| 622 | def __getattr__(self, name): |
| 623 | return getattr(self.model[0], name) |
| 624 | |
| 625 | @torch.autocast('cuda', dtype=AUTOCAST_DTYPE) |
| 626 | def forward(self, inputs): |
| 627 | x_B_T_H_W_D, t_embedding_B_T_D, crossattn_emb, rope_emb_L_1_1_D, adaln_lora_B_T_3D, timesteps_B_T = inputs |
| 628 | x_B_T_H_W_O = self.final_layer(x_B_T_H_W_D, t_embedding_B_T_D, adaln_lora_B_T_3D=adaln_lora_B_T_3D) |
| 629 | net_output_B_C_T_H_W = self.unpatchify(x_B_T_H_W_O) |
| 630 | return net_output_B_C_T_H_W |