VideoGen/encode.py

from __future__ import annotations

import shutil
import subprocess
from dataclasses import dataclass
from pathlib import Path

from PIL import Image


@dataclass
class EncodeResult:
    path: Path
    encoder: str
    mp4: bool
    gif_path: Path | None = None
    audio_path: Path | None = None
    muxed_audio: bool = False


def ffmpeg_available() -> bool:
    return shutil.which("ffmpeg") is not None


def encode_video(frames: list[Image.Image], output_path: str | Path, *, fps: int, export_gif: bool = False) -> EncodeResult:
    path = Path(output_path)
    path.parent.mkdir(parents=True, exist_ok=True)
    fps = max(1, int(fps))
    if path.suffix.lower() == ".gif":
        frames[0].save(path, save_all=True, append_images=frames[1:], duration=int(1000 / fps), loop=0)
        return EncodeResult(path, "pil_gif", False, path)
    try:
        import imageio.v3 as iio
        import numpy as np

        iio.imwrite(path, [np.asarray(frame.convert("RGB")) for frame in frames], fps=fps)
        gif_path = _write_gif(frames, path.with_suffix(".gif"), fps) if export_gif else None
        return EncodeResult(path, "imageio", path.suffix.lower() == ".mp4", gif_path)
    except Exception:
        pass
    ffmpeg = shutil.which("ffmpeg")
    if ffmpeg:
        tmp = path.parent / f"{path.stem}_frames"
        tmp.mkdir(parents=True, exist_ok=True)
        for idx, frame in enumerate(frames):
            frame.save(tmp / f"frame_{idx:05d}.png")
        subprocess.run(
            [
                ffmpeg,
                "-y",
                "-framerate",
                str(fps),
                "-i",
                str(tmp / "frame_%05d.png"),
                "-pix_fmt",
                "yuv420p",
                str(path),
            ],
            check=True,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
        gif_path = _write_gif(frames, path.with_suffix(".gif"), fps) if export_gif else None
        return EncodeResult(path, "ffmpeg", path.suffix.lower() == ".mp4", gif_path)
    fallback = path.with_suffix(".gif")
    _write_gif(frames, fallback, fps)
    return EncodeResult(fallback, "pil_gif_fallback", False, fallback)


def _write_gif(frames: list[Image.Image], path: Path, fps: int) -> Path:
    frames[0].save(path, save_all=True, append_images=frames[1:], duration=int(1000 / fps), loop=0)
    return path


def mux_audio(video_path: str | Path, audio_path: str | Path, output_path: str | Path | None = None) -> Path | None:
    ffmpeg = shutil.which("ffmpeg")
    if not ffmpeg:
        return None
    video = Path(video_path)
    audio = Path(audio_path)
    if not video.exists() or not audio.exists():
        return None
    out = Path(output_path) if output_path is not None else video.with_name(f"{video.stem}_with_audio{video.suffix}")
    out.parent.mkdir(parents=True, exist_ok=True)
    command = [
        ffmpeg,
        "-y",
        "-i",
        str(video),
        "-i",
        str(audio),
        "-map",
        "0:v:0",
        "-map",
        "1:a:0",
        "-c:v",
        "copy",
        "-c:a",
        "aac",
        "-shortest",
        str(out),
    ]
    subprocess.run(command, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    return out