(
self,
x,
t,
context,
seq_len,
clip_fea=None,
y=None,
audio=None,
ref_target_masks=None,
)
| 596 | self.enable_teacache = False |
| 597 | |
| 598 | def forward( |
| 599 | self, |
| 600 | x, |
| 601 | t, |
| 602 | context, |
| 603 | seq_len, |
| 604 | clip_fea=None, |
| 605 | y=None, |
| 606 | audio=None, |
| 607 | ref_target_masks=None, |
| 608 | ): |
| 609 | assert clip_fea is not None and y is not None |
| 610 | |
| 611 | _, T, H, W = x[0].shape |
| 612 | N_t = T // self.patch_size[0] |
| 613 | N_h = H // self.patch_size[1] |
| 614 | N_w = W // self.patch_size[2] |
| 615 | |
| 616 | if y is not None: |
| 617 | x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)] |
| 618 | x[0] = x[0].to(context[0].dtype) |
| 619 | |
| 620 | # embeddings |
| 621 | x = [self.patch_embedding(u.unsqueeze(0)) for u in x] |
| 622 | grid_sizes = torch.stack( |
| 623 | [torch.tensor(u.shape[2:], dtype=torch.long) for u in x]) |
| 624 | x = [u.flatten(2).transpose(1, 2) for u in x] |
| 625 | seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long) |
| 626 | assert seq_lens.max() <= seq_len |
| 627 | x = torch.cat([ |
| 628 | torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], |
| 629 | dim=1) for u in x |
| 630 | ]) |
| 631 | |
| 632 | # time embeddings |
| 633 | with amp.autocast(dtype=torch.float32): |
| 634 | e = self.time_embedding( |
| 635 | sinusoidal_embedding_1d(self.freq_dim, t).float()) |
| 636 | e0 = self.time_projection(e).unflatten(1, (6, self.dim)) |
| 637 | assert e.dtype == torch.float32 and e0.dtype == torch.float32 |
| 638 | |
| 639 | # text embedding |
| 640 | context_lens = None |
| 641 | context = self.text_embedding( |
| 642 | torch.stack([ |
| 643 | torch.cat( |
| 644 | [u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) |
| 645 | for u in context |
| 646 | ])) |
| 647 | |
| 648 | # clip embedding |
| 649 | if clip_fea is not None: |
| 650 | context_clip = self.img_emb(clip_fea) |
| 651 | context = torch.concat([context_clip, context], dim=1).to(x.dtype) |
| 652 | |
| 653 | |
| 654 | audio_cond = audio.to(device=x.device, dtype=x.dtype) |
| 655 | first_frame_audio_emb_s = audio_cond[:, :1, ...] |
nothing calls this directly
no test coverage detected