(
self,
video_latent,
shift,
scale,
aug_shift,
aug_scale,
cond_inds,
pred_inds
)
| 610 | ) |
| 611 | |
| 612 | def separate_modulate( |
| 613 | self, |
| 614 | video_latent, |
| 615 | shift, |
| 616 | scale, |
| 617 | aug_shift, |
| 618 | aug_scale, |
| 619 | cond_inds, |
| 620 | pred_inds |
| 621 | ): |
| 622 | video_latent = rearrange(video_latent, 'b (t n) d -> b t n d', t=self.compressed_num_frames) |
| 623 | cond_input, pred_input = video_latent[:, cond_inds], video_latent[:, pred_inds] |
| 624 | cond_input, pred_input = map( |
| 625 | lambda x: rearrange(x, 'b t n d -> b (t n) d'), |
| 626 | (cond_input, pred_input) |
| 627 | ) |
| 628 | cond_input = modulate(cond_input, aug_shift, aug_scale) |
| 629 | pred_input = modulate(pred_input, shift, scale) |
| 630 | |
| 631 | cond_input = rearrange(cond_input, 'b (t n) d -> b t n d', t=len(cond_inds)) |
| 632 | pred_input = rearrange(pred_input, 'b (t n) d -> b t n d', t=len(pred_inds)) |
| 633 | |
| 634 | video_latent = torch.cat([cond_input, pred_input], dim=1) # (b, t, n, d) |
| 635 | video_latent = rearrange(video_latent, 'b t n d -> b (t n) d') |
| 636 | return video_latent |
| 637 | |
| 638 | def separate_gating( |
| 639 | self, |
no test coverage detected