utils/video_utils.py

"""Video utilities — encode SVD frames to MP4."""
from __future__ import annotations

from pathlib import Path
from typing import List, Union

import imageio.v2 as imageio
import numpy as np
from PIL import Image


def frames_to_mp4(frames: List[Image.Image], path: Union[str, Path], fps: int = 8) -> None:
    """Save a list of PIL frames as an MP4 video."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    arr = [np.array(f.convert("RGB")) for f in frames]
    writer = imageio.get_writer(str(path), fps=fps, codec="libx264",
                                quality=8, macro_block_size=1)
    for frame in arr:
        writer.append_data(frame)
    writer.close()


def gif_from_frames(frames: List[Image.Image], path: Union[str, Path], fps: int = 8) -> None:
    """Save frames as an animated GIF (good for previews)."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    frames[0].save(
        path, save_all=True, append_images=frames[1:],
        duration=int(1000 / fps), loop=0, optimize=True,
    )