Save as frames? about animatediff HOT 1 OPEN

def save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=6, fps=8):

    # make frames dir
    frame_dir = os.path.join(os.path.dirname(path), 'frames') 
    os.makedirs(frame_dir, exist_ok=True)

    videos = rearrange(videos, "b c t h w -> t b c h w")
    os.makedirs(os.path.dirname(path), exist_ok=True)
    video_output = path.replace(".gif", ".mp4")
    writer = imageio.get_writer(video_output, fps=fps)
    outputs = []

    for i, x in enumerate(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  
        x = (x * 255).numpy().astype(np.uint8)
        
        # saved frames as PNG
        img_path = os.path.join(frame_dir, '%06d.png' % i) 
        imageio.imwrite(img_path, x)
        
        outputs.append(x)
        writer.append_data(x)

    imageio.mimsave(path, outputs, fps=fps)
    writer.close()