Spaces:
Running
on
T4
Running
on
T4
| import os | |
| import imageio | |
| import numpy as np | |
| import torch | |
| import torchvision | |
| from einops import rearrange | |
| def save_videos_grid( | |
| videos: torch.Tensor, | |
| save_path: str = 'output', | |
| path: str = 'output.gif', | |
| rescale=False, | |
| n_rows=4, | |
| duration=3 | |
| ): | |
| videos = rearrange(videos, "b c t h w -> t b c h w") | |
| outputs = [] | |
| for x in videos: | |
| x = torchvision.utils.make_grid(x, nrow=n_rows) | |
| x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) | |
| if rescale: | |
| x = (x + 1.0) / 2.0 # -1,1 -> 0,1 | |
| x = (x * 255).numpy().astype(np.uint8) | |
| outputs.append(x) | |
| if not os.path.exists(save_path): | |
| os.makedirs(save_path) | |
| imageio.mimsave(os.path.join(save_path, path), outputs, duration=duration) | |
| return os.path.join(save_path, path) |