(self, logits, **kwargs)
| 510 | return video_latent |
| 511 | |
| 512 | def final_forward(self, logits, **kwargs): |
| 513 | x, emb = logits[:, kwargs["text_length"] :, :], kwargs["emb"] # x:(b,(t n),d) |
| 514 | |
| 515 | split_cond_flag = kwargs['split_cond_flag'] |
| 516 | shift, scale = self.adaLN_modulation(emb).chunk(2, dim=1) |
| 517 | |
| 518 | if split_cond_flag: |
| 519 | cond_inds = kwargs["cond_inds"] |
| 520 | pred_inds = [i for i in range(self.compressed_num_frames) if i not in cond_inds] |
| 521 | assert cond_inds[-1] < pred_inds[0], f"cond frames must be ahead of pred frames, {cond_inds}, {pred_inds}" |
| 522 | aug_emb = kwargs["aug_emb"] |
| 523 | aug_shift, aug_scale = self.adaLN_modulation(aug_emb).chunk(2, dim=1) |
| 524 | |
| 525 | if split_cond_flag: |
| 526 | x = self.separate_modulate(self.norm_final(x), shift, scale, aug_shift, aug_scale, cond_inds, pred_inds) |
| 527 | else: |
| 528 | x = modulate(self.norm_final(x), shift, scale) |
| 529 | x = self.linear(x) |
| 530 | |
| 531 | return unpatchify( |
| 532 | x, |
| 533 | c=self.out_channels, |
| 534 | p=self.patch_size, |
| 535 | w=self.latent_width // self.patch_size, |
| 536 | h=self.latent_height // self.patch_size, |
| 537 | rope_position_ids=kwargs.get("rope_position_ids", None), |
| 538 | **kwargs, |
| 539 | ) |
| 540 | |
| 541 | def reinit(self, parent_model=None): |
| 542 | nn.init.xavier_uniform_(self.linear.weight) |
nothing calls this directly
no test coverage detected