wehub-resource-sync
245 lines
8.0 KiB
Python
245 lines
8.0 KiB
Python
# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
"""Progressive MP4 writer for the chunk-pipelined Sana-WM streaming inference.
|
|
|
|
Spawns a single ``ffmpeg`` subprocess that consumes raw RGB frames via stdin
|
|
and emits a growing H.264 MP4 on disk. Designed for the streaming pipeline
|
|
where the orchestrator pushes one (T, H, W, 3) uint8 chunk per AR block as
|
|
decoding completes.
|
|
|
|
The writer keeps no buffering of its own — frames are forwarded byte-for-byte
|
|
to ffmpeg. Concurrency model: the orchestrator drives ``write_chunk`` from the
|
|
host thread after a CUDA-stream ``synchronize`` ensures the decoded pixels
|
|
have been copied to CPU. ``close`` blocks until ffmpeg has flushed and
|
|
finalized the MOOV atom so the resulting file is immediately playable.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import shlex
|
|
import shutil
|
|
import subprocess
|
|
from pathlib import Path
|
|
from typing import IO
|
|
|
|
import numpy as np
|
|
|
|
|
|
def resolve_ffmpeg_exe() -> str:
|
|
"""Return an ffmpeg executable path, falling back to imageio-ffmpeg."""
|
|
env_path = os.environ.get("SANA_WM_FFMPEG_BIN", "").strip()
|
|
if env_path:
|
|
return env_path
|
|
system_path = shutil.which("ffmpeg")
|
|
if system_path:
|
|
return system_path
|
|
try:
|
|
import imageio_ffmpeg
|
|
|
|
return imageio_ffmpeg.get_ffmpeg_exe()
|
|
except Exception as exc:
|
|
raise RuntimeError(
|
|
"ffmpeg was not found on PATH and imageio_ffmpeg is unavailable. "
|
|
"Install ffmpeg or set SANA_WM_FFMPEG_BIN."
|
|
) from exc
|
|
|
|
|
|
class StreamingMp4Writer:
|
|
"""Pipe raw RGB chunks to a single ffmpeg process; produce H.264 MP4.
|
|
|
|
Args:
|
|
output_path: Destination ``.mp4`` path. Parent directory is created.
|
|
height: Pixel height of every frame.
|
|
width: Pixel width of every frame.
|
|
fps: Output frame rate.
|
|
crf: ffmpeg constant rate factor (lower = higher quality; 18 is
|
|
near-lossless visually).
|
|
preset: ffmpeg libx264 preset (``slow``/``medium``/``fast``/...).
|
|
extra_args: Optional list of extra ffmpeg CLI arguments inserted
|
|
between ``-i pipe:0`` and the encoder flags.
|
|
loglevel: ffmpeg log level (default ``warning``).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
output_path: str | Path,
|
|
*,
|
|
height: int,
|
|
width: int,
|
|
fps: int = 16,
|
|
crf: int = 18,
|
|
preset: str = "medium",
|
|
encoder: str = "libx264",
|
|
extra_args: list[str] | None = None,
|
|
loglevel: str = "warning",
|
|
) -> None:
|
|
self._path = Path(output_path).expanduser().resolve()
|
|
self._path.parent.mkdir(parents=True, exist_ok=True)
|
|
self._H = int(height)
|
|
self._W = int(width)
|
|
self._fps = int(fps)
|
|
self._closed = False
|
|
self._frames_written = 0
|
|
|
|
encoder = str(encoder).strip().lower()
|
|
if encoder in {"x264", "cpu", "libx264"}:
|
|
codec_args = [
|
|
"-c:v",
|
|
"libx264",
|
|
"-preset",
|
|
preset,
|
|
"-crf",
|
|
str(int(crf)),
|
|
]
|
|
elif encoder in {"nvenc", "h264_nvenc"}:
|
|
# Use NVIDIA's hardware H.264 encoder. The raw RGB upload and
|
|
# yuv420 conversion remain part of the ffmpeg process, but entropy
|
|
# coding no longer competes with the Python process on CPU cores.
|
|
codec_args = [
|
|
"-c:v",
|
|
"h264_nvenc",
|
|
"-preset",
|
|
preset,
|
|
"-cq",
|
|
str(int(crf)),
|
|
"-b:v",
|
|
"0",
|
|
]
|
|
else:
|
|
raise ValueError(f"Unsupported streaming MP4 encoder: {encoder!r}.")
|
|
|
|
cmd: list[str] = [
|
|
resolve_ffmpeg_exe(),
|
|
"-y",
|
|
"-loglevel",
|
|
loglevel,
|
|
"-f",
|
|
"rawvideo",
|
|
"-vcodec",
|
|
"rawvideo",
|
|
"-pix_fmt",
|
|
"rgb24",
|
|
"-s",
|
|
f"{self._W}x{self._H}",
|
|
"-r",
|
|
str(self._fps),
|
|
"-i",
|
|
"pipe:0",
|
|
*(extra_args or []),
|
|
*codec_args,
|
|
"-pix_fmt",
|
|
"yuv420p",
|
|
"-movflags",
|
|
"+faststart",
|
|
"-r",
|
|
str(self._fps),
|
|
str(self._path),
|
|
]
|
|
self._cmd = cmd
|
|
self._proc = subprocess.Popen(
|
|
cmd,
|
|
stdin=subprocess.PIPE,
|
|
stdout=subprocess.DEVNULL,
|
|
stderr=subprocess.PIPE,
|
|
bufsize=0,
|
|
)
|
|
|
|
@property
|
|
def output_path(self) -> Path:
|
|
return self._path
|
|
|
|
@property
|
|
def frames_written(self) -> int:
|
|
return self._frames_written
|
|
|
|
@property
|
|
def ffmpeg_command(self) -> str:
|
|
return shlex.join(self._cmd)
|
|
|
|
def write_chunk(self, frames_uint8: np.ndarray) -> None:
|
|
"""Append a chunk of frames.
|
|
|
|
Args:
|
|
frames_uint8: ``(T, H, W, 3)`` uint8 RGB tensor with ``H == height``
|
|
and ``W == width`` as configured at construction time.
|
|
"""
|
|
if self._closed:
|
|
raise RuntimeError("write_chunk called after close().")
|
|
if frames_uint8.dtype != np.uint8:
|
|
raise ValueError(f"frames must be uint8; got dtype {frames_uint8.dtype}.")
|
|
if frames_uint8.ndim != 4 or frames_uint8.shape[-1] != 3:
|
|
raise ValueError(f"frames must have shape (T,H,W,3); got {frames_uint8.shape}.")
|
|
if frames_uint8.shape[1] != self._H or frames_uint8.shape[2] != self._W:
|
|
raise ValueError(f"frame H,W = {frames_uint8.shape[1:3]} but writer expects {(self._H, self._W)}.")
|
|
if not frames_uint8.flags["C_CONTIGUOUS"]:
|
|
frames_uint8 = np.ascontiguousarray(frames_uint8)
|
|
|
|
stdin: IO[bytes] | None = self._proc.stdin
|
|
if stdin is None:
|
|
raise RuntimeError("ffmpeg stdin is None; subprocess failed to start.")
|
|
try:
|
|
stdin.write(frames_uint8.tobytes())
|
|
except BrokenPipeError as exc:
|
|
stderr_blob = b""
|
|
if self._proc.stderr is not None:
|
|
stderr_blob = self._proc.stderr.read() or b""
|
|
raise RuntimeError(
|
|
f"ffmpeg stdin BrokenPipeError; ffmpeg likely exited.\n"
|
|
f"command: {self.ffmpeg_command}\n"
|
|
f"stderr:\n{stderr_blob.decode(errors='replace')}"
|
|
) from exc
|
|
self._frames_written += int(frames_uint8.shape[0])
|
|
|
|
def close(self) -> Path:
|
|
"""Flush stdin, wait for ffmpeg to finalize, and return the output path.
|
|
|
|
Idempotent — calling twice is a no-op on the second call. Raises
|
|
``RuntimeError`` if ffmpeg returned a non-zero exit code, including
|
|
the captured stderr for diagnostics.
|
|
"""
|
|
if self._closed:
|
|
return self._path
|
|
self._closed = True
|
|
if self._proc.stdin is not None:
|
|
try:
|
|
self._proc.stdin.close()
|
|
except BrokenPipeError:
|
|
pass
|
|
stderr_blob = b""
|
|
if self._proc.stderr is not None:
|
|
stderr_blob = self._proc.stderr.read() or b""
|
|
rc = self._proc.wait()
|
|
if rc != 0:
|
|
raise RuntimeError(
|
|
f"ffmpeg exited with code {rc}\n"
|
|
f"command: {self.ffmpeg_command}\n"
|
|
f"stderr:\n{stderr_blob.decode(errors='replace')}"
|
|
)
|
|
return self._path
|
|
|
|
def __enter__(self) -> StreamingMp4Writer:
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc, tb) -> None:
|
|
try:
|
|
if exc_type is None:
|
|
self.close()
|
|
else:
|
|
try:
|
|
if self._proc.stdin is not None:
|
|
self._proc.stdin.close()
|
|
except Exception:
|
|
pass
|
|
self._proc.terminate()
|
|
finally:
|
|
pass
|