| 42 | |
| 43 | |
| 44 | def save_video_as_grid_and_mp4(video_batch: torch.Tensor, save_path: str, T: int, fps: int = 5, args=None, key=None): |
| 45 | os.makedirs(save_path, exist_ok=True) |
| 46 | |
| 47 | for i, vid in enumerate(video_batch): |
| 48 | gif_frames = [] |
| 49 | for frame in vid: |
| 50 | frame = rearrange(frame, "c h w -> h w c") |
| 51 | frame = (255.0 * frame).cpu().numpy().astype(np.uint8) |
| 52 | gif_frames.append(frame) |
| 53 | now_save_path = os.path.join(save_path, f"{i:06d}.mp4") |
| 54 | with imageio.get_writer(now_save_path, fps=fps) as writer: |
| 55 | for frame in gif_frames: |
| 56 | writer.append_data(frame) |
| 57 | if args is not None and args.wandb: |
| 58 | wandb.log( |
| 59 | {key + f"_video_{i}": wandb.Video(now_save_path, fps=fps, format="mp4")}, step=args.iteration + 1 |
| 60 | ) |
| 61 | |
| 62 | |
| 63 | def log_video(batch, model, args, only_log_video_latents=False): |