| 46 | return torch.cat(samples, dim=1) |
| 47 | |
| 48 | def sample_loop(self, xT, x_cond, motion_frames, audio_emb): |
| 49 | xt = xT |
| 50 | for i, t in reversed(list(enumerate(self.timesteps))): |
| 51 | timesteps = torch.tensor([t] * xT.shape[0]).to(xT.device) |
| 52 | timesteps_ids = torch.tensor([i] * xT.shape[0]).to(xT.device) |
| 53 | nn_out = self.nn_backbone(xt, timesteps, x_cond, motion_frames=motion_frames, audio_emb=audio_emb) |
| 54 | mean, logvar = self.get_p_params(xt, timesteps_ids, nn_out) |
| 55 | noise = torch.randn_like(xt) if t > 0 else torch.zeros_like(xt) |
| 56 | xt = mean + noise * torch.exp(logvar / 2) |
| 57 | |
| 58 | return xt |
| 59 | |
| 60 | def get_p_params(self, xt, timesteps, nn_out): |
| 61 | if self.in_channels == self.out_channels: |