(
self,
video_latent,
shift,
scale,
aug_shift,
aug_scale,
cond_inds,
pred_inds
)
| 484 | self.compressed_num_frames = compressed_num_frames |
| 485 | |
| 486 | def separate_modulate( |
| 487 | self, |
| 488 | video_latent, |
| 489 | shift, |
| 490 | scale, |
| 491 | aug_shift, |
| 492 | aug_scale, |
| 493 | cond_inds, |
| 494 | pred_inds |
| 495 | ): |
| 496 | video_latent = rearrange(video_latent, 'b (t n) d -> b t n d', t=self.compressed_num_frames) |
| 497 | cond_input, pred_input = video_latent[:, cond_inds], video_latent[:, pred_inds] |
| 498 | cond_input, pred_input = map( |
| 499 | lambda x: rearrange(x, 'b t n d -> b (t n) d'), |
| 500 | (cond_input, pred_input) |
| 501 | ) |
| 502 | cond_input = modulate(cond_input, aug_shift, aug_scale) |
| 503 | pred_input = modulate(pred_input, shift, scale) |
| 504 | |
| 505 | cond_input = rearrange(cond_input, 'b (t n) d -> b t n d', t=len(cond_inds)) |
| 506 | pred_input = rearrange(pred_input, 'b (t n) d -> b t n d', t=len(pred_inds)) |
| 507 | |
| 508 | video_latent = torch.cat([cond_input, pred_input], dim=1) # (b, t, n, d) |
| 509 | video_latent = rearrange(video_latent, 'b t n d -> b (t n) d') |
| 510 | return video_latent |
| 511 | |
| 512 | def final_forward(self, logits, **kwargs): |
| 513 | x, emb = logits[:, kwargs["text_length"] :, :], kwargs["emb"] # x:(b,(t n),d) |
no test coverage detected