(video, frame_indices)
| 49 | |
| 50 | |
| 51 | def pick_video_frame(video, frame_indices): |
| 52 | batch, device = video.shape[0], video.device |
| 53 | video = rearrange(video, "b c f ... -> b f c ...") |
| 54 | batch_indices = torch.arange(batch, device=device) |
| 55 | batch_indices = rearrange(batch_indices, "b -> b 1") |
| 56 | images = video[batch_indices, frame_indices] |
| 57 | images = rearrange(images, "b 1 c ... -> b c ...") |
| 58 | return images |
| 59 | |
| 60 | |
| 61 | def gradient_penalty(images, output): |