scripts/media_packaging.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
from contextlib import contextmanager
import html
import os
import subprocess
from pathlib import Path
from typing import Iterator
import uuid


@contextmanager
def atomic_media_output(output: Path) -> Iterator[Path]:
    output.parent.mkdir(parents=True, exist_ok=True)
    partial = output.with_name(
        f"{output.stem}.partial-{os.getpid()}-{uuid.uuid4().hex}{output.suffix}"
    )
    try:
        yield partial
        with partial.open("rb") as handle:
            os.fsync(handle.fileno())
        os.replace(partial, output)
    finally:
        partial.unlink(missing_ok=True)


def hevc_command(source: Path, output: Path) -> list[str]:
    return [
        "ffmpeg",
        "-y",
        "-i",
        str(source),
        "-map",
        "0:v:0",
        "-map",
        "0:a:0?",
        "-c:v",
        "libx265",
        "-tag:v",
        "hvc1",
        "-pix_fmt",
        "yuv420p",
        "-preset",
        "medium",
        "-crf",
        "24",
        "-c:a",
        "copy",
        "-movflags",
        "+faststart",
        str(output),
    ]


def h264_command(source: Path, output: Path) -> list[str]:
    return [
        "ffmpeg",
        "-y",
        "-i",
        str(source),
        "-map",
        "0:v:0",
        "-map",
        "0:a:0?",
        "-c:v",
        "libx264",
        "-pix_fmt",
        "yuv420p",
        "-preset",
        "medium",
        "-crf",
        "20",
        "-c:a",
        "copy",
        "-movflags",
        "+faststart",
        str(output),
    ]


def embedded_video_html(*, repo_id: str, stem: str, title: str) -> str:
    base = f"https://huggingface.co/{repo_id}/resolve/main/examples"
    label = html.escape(title, quote=True)
    hevc = html.escape(f"{base}/hevc/{stem}.mp4", quote=True)
    h264 = html.escape(f"{base}/h264/{stem}.mp4", quote=True)
    return (
        f'<video controls playsinline preload="metadata" width="608" aria-label="{label}">\n'
        f'  <source src="{hevc}" type="video/mp4; codecs=&quot;hvc1&quot;">\n'
        f'  <source src="{h264}" type="video/mp4">\n'
        f'  <a href="{h264}">{label} MP4</a>\n'
        "</video>"
    )


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--source", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    with atomic_media_output(args.output) as partial:
        subprocess.run(hevc_command(args.source, partial), check=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())