| 531 | |
| 532 | class InitialLayer(nn.Module): |
| 533 | def __init__(self, model, text_encoder, is_generic_llm): |
| 534 | super().__init__() |
| 535 | self.x_embedder = model.x_embedder |
| 536 | self.pos_embedder = model.pos_embedder |
| 537 | if model.extra_per_block_abs_pos_emb: |
| 538 | self.extra_pos_embedder = model.extra_pos_embedder |
| 539 | self.t_embedder = model.t_embedder |
| 540 | self.t_embedding_norm = model.t_embedding_norm |
| 541 | self.text_encoder = text_encoder |
| 542 | self.model = [model] |
| 543 | self.is_generic_llm = is_generic_llm |
| 544 | |
| 545 | @torch.autocast('cuda', dtype=AUTOCAST_DTYPE) |
| 546 | def forward(self, inputs): |