x: motion latents of shape [B, T, C]. x_mask: mask of shape [B, T], 1 for valid, 0 for invalid. t: [B]. context: text features of shape [B, L, C]. ref_motion: [B, T, C]. ref_motion_mask: [B, T], 1 for valid
(
self,
x,
timestep,
context,
clip_fea=None,
y=None,
use_gradient_checkpointing=False,
x_mask=None,
ref_motion=None,
ref_motion_mask=None,
attend_to_text_mask=None,
)
| 648 | return list(self.blocks) |
| 649 | |
| 650 | def forward( |
| 651 | self, |
| 652 | x, |
| 653 | timestep, |
| 654 | context, |
| 655 | clip_fea=None, |
| 656 | y=None, |
| 657 | use_gradient_checkpointing=False, |
| 658 | x_mask=None, |
| 659 | ref_motion=None, |
| 660 | ref_motion_mask=None, |
| 661 | attend_to_text_mask=None, |
| 662 | ): |
| 663 | """ |
| 664 | x: motion latents of shape [B, T, C]. |
| 665 | x_mask: mask of shape [B, T], 1 for valid, 0 for invalid. |
| 666 | t: [B]. |
| 667 | context: text features of shape [B, L, C]. |
| 668 | ref_motion: [B, T, C]. |
| 669 | ref_motion_mask: [B, T], 1 for valid, 0 for invalid. |
| 670 | attend_to_text_mask: [B], 1 for attend to text, 0 for attend to ref motion. |
| 671 | """ |
| 672 | def create_custom_forward(module): |
| 673 | def custom_forward(*inputs, **kwargs): |
| 674 | return module(*inputs, **kwargs) |
| 675 | return custom_forward |
| 676 | |
| 677 | if attend_to_text_mask is None: |
| 678 | # if ref motion is not available, it will be set as all zeros tensor. |
| 679 | attend_to_text_mask = ((ref_motion**2).sum(dim=[1,2]) == 0) |
| 680 | # print('computed attend_to_text_mask:', attend_to_text_mask) |
| 681 | |
| 682 | if self.model_type == 'i2v': |
| 683 | assert clip_fea is not None and y is not None |
| 684 | # params |
| 685 | device = x[0].device |
| 686 | if self.freqs.device != device: |
| 687 | self.freqs = self.freqs.to(device) |
| 688 | |
| 689 | # embeddings |
| 690 | x = self.motion_embedding(x) |
| 691 | x = self.motion_pos_embedding(x) |
| 692 | |
| 693 | # ref motion |
| 694 | if ref_motion is not None: |
| 695 | ref_motion = self.ref_motion_embedding(ref_motion) |
| 696 | ref_motion = self.ref_motion_pos_embedding(ref_motion) |
| 697 | |
| 698 | # get seq_lens based on x_mask |
| 699 | if x_mask is None: |
| 700 | x_mask = torch.ones((x.size(0), x.size(1)), device=x.device) |
| 701 | seq_lens = x_mask.sum(1).to(torch.long) |
| 702 | |
| 703 | # time embeddings |
| 704 | with amp.autocast(dtype=torch.float32, device_type="cuda"): |
| 705 | e = self.time_embedding( |
| 706 | sinusoidal_embedding_1d(self.freq_dim, timestep).float()) |
| 707 | e0 = self.time_projection(e).unflatten(1, (6, self.dim)) |
nothing calls this directly
no test coverage detected