| 25 | self.device = device |
| 26 | |
| 27 | def sample(self, x_cond, audio_emb, n_audio_motion_embs=2, n_motion_frames=2, motion_channels=3): |
| 28 | with torch.no_grad(): |
| 29 | n_frames = audio_emb.shape[1] |
| 30 | |
| 31 | xT = torch.randn(x_cond.shape[0], n_frames, self.in_channels, self.x_shape[0], self.x_shape[1]).to(x_cond.device) |
| 32 | |
| 33 | audio_ids = [0] * n_audio_motion_embs |
| 34 | for i in range(n_audio_motion_embs + 1): |
| 35 | audio_ids += [i] |
| 36 | |
| 37 | motion_frames = [self.motion_transforms(x_cond) for _ in range(n_motion_frames)] |
| 38 | motion_frames = torch.cat(motion_frames, dim=1) |
| 39 | |
| 40 | samples = [] |
| 41 | for i in trange(n_frames, desc=f'Sampling'): |
| 42 | sample_frame = self.sample_loop(xT[:, i].to(x_cond.device), x_cond, motion_frames, audio_emb[:, audio_ids]) |
| 43 | samples.append(sample_frame.unsqueeze(1)) |
| 44 | motion_frames = torch.cat([motion_frames[:, motion_channels:, :], self.motion_transforms(sample_frame)], dim=1) |
| 45 | audio_ids = audio_ids[1:] + [min(i + n_audio_motion_embs + 1, n_frames - 1)] |
| 46 | return torch.cat(samples, dim=1) |
| 47 | |
| 48 | def sample_loop(self, xT, x_cond, motion_frames, audio_emb): |
| 49 | xt = xT |