501 lines
18 KiB
Python
501 lines
18 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||
"""Video/audio muxing utilities using PyAV (no ffmpeg binary dependency)."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import io
|
||
import queue
|
||
import threading
|
||
from collections.abc import Callable, Iterable, Mapping
|
||
from fractions import Fraction
|
||
from typing import Any, NamedTuple, cast
|
||
|
||
import av
|
||
import numpy as np
|
||
from vllm.logger import init_logger
|
||
|
||
logger = init_logger(__name__)
|
||
|
||
_CHUNKED_MP4_DONE = object()
|
||
|
||
|
||
class _QueuedChunk(NamedTuple):
|
||
"""One queued chunk plus the callback that returns its buffer to its owner."""
|
||
|
||
frames: np.ndarray
|
||
on_consumed: Callable[[], None] | None
|
||
|
||
|
||
def _release_chunk(on_consumed: Callable[[], None] | None) -> None:
|
||
"""Hand a chunk's buffer back, never letting the callback wedge the worker."""
|
||
if on_consumed is None:
|
||
return
|
||
try:
|
||
on_consumed()
|
||
except BaseException: # noqa: BLE001
|
||
logger.exception("Chunked MP4 encoder failed to release a consumed chunk")
|
||
|
||
|
||
def normalize_preencode_batch_frames(value: Any) -> int:
|
||
"""Validate the request's MP4 transfer/encoding batch threshold."""
|
||
if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
|
||
raise ValueError("preencode_batch_frames must be a positive integer")
|
||
return value
|
||
|
||
|
||
def normalize_video_codec_options(value: Any) -> dict[str, str] | None:
|
||
"""Coerce a request's ``video_codec_options`` into PyAV's str->str contract.
|
||
|
||
The value reaches a worker straight from ``extra_params``, so reject a
|
||
non-mapping here rather than letting it fail inside the encoder after the
|
||
decode has already run.
|
||
"""
|
||
if value is None:
|
||
return None
|
||
if not isinstance(value, Mapping):
|
||
raise ValueError("video_codec_options must be a JSON object of encoder option names to values")
|
||
return {str(key): str(item) for key, item in value.items()}
|
||
|
||
|
||
def _validate_video_chunk(chunk: np.ndarray, *, width: int, height: int) -> None:
|
||
"""Validate a ``(T, H, W, 3)`` uint8 RGB chunk against a session's frame size."""
|
||
if chunk.ndim != 4 or chunk.shape[-1] != 3:
|
||
raise ValueError("video chunk must have shape (T, H, W, 3)")
|
||
if chunk.dtype != np.uint8:
|
||
raise ValueError("video chunk must have dtype uint8")
|
||
if chunk.shape[1] != height or chunk.shape[2] != width:
|
||
raise ValueError("video chunks in a session must use a consistent frame size")
|
||
|
||
|
||
class ChunkedMP4Encoder:
|
||
"""Encode temporal video chunks while the producer is still decoding.
|
||
|
||
A bounded queue and one muxing worker provide ordered backpressure. What
|
||
that bounds is pending raw-frame memory: at most ``max_pending`` chunks wait
|
||
to be encoded, so the bound is in chunks rather than bytes and a taller
|
||
frame or a longer chunk raises it. The muxed container is not bounded -- it
|
||
accumulates in memory and grows with output duration. Chunks use the same
|
||
``(T, H, W, 3)`` uint8 contract as :func:`mux_video_audio_bytes`.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
*,
|
||
width: int,
|
||
height: int,
|
||
fps: float,
|
||
audio_waveform: np.ndarray | None = None,
|
||
audio_sample_rate: int | None = None,
|
||
max_pending: int = 2,
|
||
video_codec: str = "h264",
|
||
audio_codec: str = "aac",
|
||
crf: str = "18",
|
||
video_codec_options: dict[str, str] | None = None,
|
||
) -> None:
|
||
if max_pending <= 0:
|
||
raise ValueError("max_pending must be positive")
|
||
if width <= 0 or height <= 0:
|
||
raise ValueError("width and height must be positive")
|
||
self.width = int(width)
|
||
self.height = int(height)
|
||
self._queue: queue.Queue[object] = queue.Queue(maxsize=max_pending)
|
||
self._result: bytes | None = None
|
||
self._error: BaseException | None = None
|
||
self._closed = False
|
||
self._aborted = False
|
||
self._input_done = False
|
||
self._state_lock = threading.Lock()
|
||
|
||
def run() -> None:
|
||
frames = self._frames()
|
||
try:
|
||
self._result = mux_av_video_audio_bytes(
|
||
frames,
|
||
width=self.width,
|
||
height=self.height,
|
||
audio_waveform=audio_waveform,
|
||
fps=fps,
|
||
audio_sample_rate=audio_sample_rate,
|
||
video_codec=video_codec,
|
||
audio_codec=audio_codec,
|
||
crf=crf,
|
||
video_codec_options=video_codec_options,
|
||
)
|
||
except BaseException as exc:
|
||
self._error = exc
|
||
# Closing is what runs the generator's cleanup and releases the
|
||
# chunk it was mid-way through. Waiting for collection would
|
||
# never get there: the traceback just stored on this encoder
|
||
# keeps the abandoned generator alive for the process's life.
|
||
frames.close()
|
||
self._drain_until_done()
|
||
|
||
self._thread = threading.Thread(target=run, name="chunked-mp4", daemon=True)
|
||
self._thread.start()
|
||
|
||
def _drain_until_done(self) -> None:
|
||
# A flush/mux failure can occur after _frames consumed the sentinel.
|
||
if self._input_done:
|
||
return
|
||
while True:
|
||
entry = self._queue.get()
|
||
if entry is _CHUNKED_MP4_DONE:
|
||
return
|
||
_release_chunk(cast(_QueuedChunk, entry).on_consumed)
|
||
|
||
def _frames(self):
|
||
while True:
|
||
entry = self._queue.get()
|
||
if entry is _CHUNKED_MP4_DONE:
|
||
self._input_done = True
|
||
return
|
||
assert isinstance(entry, _QueuedChunk)
|
||
try:
|
||
for frame_data in entry.frames:
|
||
yield av.VideoFrame.from_ndarray(frame_data, format="rgb24")
|
||
finally:
|
||
# Also runs when the muxer raises part-way through this chunk
|
||
# and abandons the generator, so a lent buffer is never
|
||
# stranded in the encoder.
|
||
_release_chunk(entry.on_consumed)
|
||
|
||
def _validate_chunk(self, chunk: np.ndarray) -> None:
|
||
_validate_video_chunk(chunk, width=self.width, height=self.height)
|
||
|
||
def _raise_if_failed(self) -> None:
|
||
if self._error is not None:
|
||
raise self._error
|
||
if self._closed:
|
||
raise RuntimeError("ChunkedMP4Encoder is already closed")
|
||
|
||
def push(self, chunk: np.ndarray, *, on_consumed: Callable[[], None] | None = None) -> None:
|
||
"""Queue a chunk; asynchronous failures surface here or at finish().
|
||
|
||
``on_consumed`` fires exactly once after this encoder has read
|
||
``chunk``, including when the chunk is discarded by an abort or by a
|
||
failed encode. Ownership transfers on call: whether this returns or
|
||
raises, the callback is this encoder's responsibility, so a caller
|
||
lending a pooled host buffer never compensates on the error path.
|
||
"""
|
||
try:
|
||
self._validate_chunk(chunk)
|
||
self._raise_if_failed()
|
||
except BaseException:
|
||
_release_chunk(on_consumed)
|
||
raise
|
||
self._queue.put(_QueuedChunk(chunk, on_consumed))
|
||
self._raise_if_failed()
|
||
|
||
def _send_done(self) -> None:
|
||
self._queue.put(_CHUNKED_MP4_DONE)
|
||
|
||
def finish(self) -> bytes:
|
||
"""Flush the muxer and return complete progressive MP4 bytes."""
|
||
with self._state_lock:
|
||
if not self._closed:
|
||
self._closed = True
|
||
self._send_done()
|
||
self._thread.join()
|
||
if self._error is not None:
|
||
raise self._error
|
||
if self._aborted:
|
||
raise RuntimeError("ChunkedMP4Encoder was aborted")
|
||
assert self._result is not None
|
||
return self._result
|
||
|
||
def abort(self) -> None:
|
||
"""Discard pending chunks and join the worker after cancellation.
|
||
|
||
Any chunk already being encoded must finish before the worker exits.
|
||
"""
|
||
with self._state_lock:
|
||
if not self._closed:
|
||
self._closed = True
|
||
self._aborted = True
|
||
while True:
|
||
try:
|
||
entry = self._queue.get_nowait()
|
||
except queue.Empty:
|
||
break
|
||
if entry is not _CHUNKED_MP4_DONE:
|
||
_release_chunk(cast(_QueuedChunk, entry).on_consumed)
|
||
self._send_done()
|
||
self._thread.join()
|
||
|
||
close = finish
|
||
|
||
def __enter__(self) -> ChunkedMP4Encoder:
|
||
return self
|
||
|
||
def __exit__(self, exc_type, exc, tb) -> None:
|
||
if exc_type is None:
|
||
self.finish()
|
||
else:
|
||
self.abort()
|
||
|
||
|
||
class FragmentedMP4Muxer:
|
||
"""Incrementally mux video frames into one fragmented MP4 byte stream."""
|
||
|
||
def __init__(
|
||
self,
|
||
*,
|
||
width: int,
|
||
height: int,
|
||
fps: float = 25.0,
|
||
video_codec: str = "h264",
|
||
crf: str = "18",
|
||
video_codec_options: dict[str, str] | None = None,
|
||
) -> None:
|
||
self._buf = io.BytesIO()
|
||
self._closed = False
|
||
self._container = av.open(
|
||
self._buf,
|
||
mode="w",
|
||
format="mp4",
|
||
options={"movflags": "+frag_every_frame+empty_moov+default_base_moof"},
|
||
)
|
||
|
||
self._stream: av.VideoStream = cast(
|
||
av.VideoStream,
|
||
self._container.add_stream(video_codec, rate=Fraction(fps).limit_denominator(10000)),
|
||
)
|
||
self._stream.width = width
|
||
self._stream.height = height
|
||
self._stream.pix_fmt = "yuv420p"
|
||
|
||
options: dict[str, object] = {"crf": str(crf)}
|
||
if video_codec_options:
|
||
options.update(video_codec_options)
|
||
self._stream.options = options
|
||
|
||
try:
|
||
self._stream.codec_context.max_b_frames = 0
|
||
except AttributeError:
|
||
pass
|
||
|
||
def mux_video_frames(self, video_frames: np.ndarray) -> bytes:
|
||
"""Mux a batch of ``uint8`` RGB frames and return newly written MP4 bytes."""
|
||
if self._closed:
|
||
raise RuntimeError("Cannot mux frames after FragmentedMP4Muxer.close().")
|
||
_validate_video_chunk(video_frames, width=self._stream.width, height=self._stream.height)
|
||
|
||
for frame_data in video_frames:
|
||
frame = av.VideoFrame.from_ndarray(frame_data, format="rgb24")
|
||
for packet in self._stream.encode(frame):
|
||
self._container.mux(packet)
|
||
return self._read_new_bytes()
|
||
|
||
def close(self) -> bytes:
|
||
"""Flush delayed encoder packets, close the container, and return final bytes."""
|
||
if self._closed:
|
||
return b""
|
||
for packet in self._stream.encode():
|
||
self._container.mux(packet)
|
||
self._container.close()
|
||
self._closed = True
|
||
return self._read_new_bytes()
|
||
|
||
def _read_new_bytes(self) -> bytes:
|
||
"""Return newly muxed bytes in the current video container,
|
||
then clear the buffer to prepare for the next chunk."""
|
||
chunk = self._buf.getvalue()
|
||
self._buf.seek(0)
|
||
self._buf.truncate()
|
||
return chunk
|
||
|
||
|
||
def count_mp4_frames(video_bytes: bytes) -> int | None:
|
||
"""Read an MP4's frame count from its sample table, without decoding it."""
|
||
if not video_bytes:
|
||
return None
|
||
try:
|
||
with cast(Any, av.open(io.BytesIO(video_bytes), format="mp4")) as container:
|
||
frames = int(container.streams.video[0].frames)
|
||
except Exception:
|
||
return None
|
||
return frames or None
|
||
|
||
|
||
def finalize_streaming_video_bytes(
|
||
video_bytes: bytes,
|
||
*,
|
||
input_format: str,
|
||
fps: float = 25.0,
|
||
video_codec_options: dict[str, str] | None = None,
|
||
) -> bytes:
|
||
"""Convert streamed video bytes into a progressive MP4 for local playback."""
|
||
if not video_bytes:
|
||
return video_bytes
|
||
|
||
normalized_format = input_format.lower()
|
||
if normalized_format == "m4s":
|
||
demux_format = "mp4"
|
||
else:
|
||
raise ValueError(f"Unsupported streaming video format: {input_format}")
|
||
|
||
try:
|
||
with cast(Any, av.open(io.BytesIO(video_bytes), format=demux_format)) as container:
|
||
stream = container.streams.video[0]
|
||
frame_arrays = [frame.to_ndarray(format="rgb24") for frame in container.decode(stream)]
|
||
except Exception:
|
||
return video_bytes
|
||
|
||
if not frame_arrays:
|
||
return video_bytes
|
||
|
||
frames_u8 = np.ascontiguousarray(np.stack(frame_arrays, axis=0), dtype=np.uint8)
|
||
return mux_video_audio_bytes(
|
||
frames_u8,
|
||
None,
|
||
fps=float(fps),
|
||
video_codec_options=video_codec_options,
|
||
)
|
||
|
||
|
||
def mux_video_audio_bytes(
|
||
video_frames: np.ndarray,
|
||
audio_waveform: np.ndarray | None = None,
|
||
*,
|
||
fps: float = 25.0,
|
||
audio_sample_rate: int = 44100,
|
||
video_codec: str = "h264",
|
||
audio_codec: str = "aac",
|
||
crf: str = "18",
|
||
video_codec_options: dict[str, str] | None = None,
|
||
) -> bytes:
|
||
"""Mux video frames and optional audio waveform into MP4 bytes.
|
||
|
||
Args:
|
||
video_frames: uint8 array of shape ``(T, H, W, 3)`` (RGB).
|
||
audio_waveform: float32 array – mono ``(N,)`` or ``(N, C)`` / ``(C, N)``.
|
||
fps: Video frame rate.
|
||
audio_sample_rate: Audio sample rate in Hz.
|
||
video_codec: Video codec name.
|
||
audio_codec: Audio codec name.
|
||
crf: Constant rate factor for the video encoder.
|
||
|
||
Returns:
|
||
Raw MP4 bytes ready to be written to disk or streamed.
|
||
"""
|
||
buf = io.BytesIO()
|
||
container = av.open(buf, mode="w", format="mp4")
|
||
|
||
v_stream = cast(av.VideoStream, container.add_stream(video_codec, rate=Fraction(fps).limit_denominator(10000)))
|
||
v_stream.width = video_frames.shape[2]
|
||
v_stream.height = video_frames.shape[1]
|
||
v_stream.pix_fmt = "yuv420p"
|
||
|
||
options: dict[str, object] = {"crf": str(crf)}
|
||
if video_codec_options:
|
||
options.update(video_codec_options)
|
||
v_stream.options = options
|
||
|
||
a_stream: av.AudioStream | None = None
|
||
samples: np.ndarray | None = None
|
||
layout: str | None = None
|
||
if audio_waveform is not None:
|
||
samples = audio_waveform.astype(np.float32)
|
||
if samples.ndim == 1:
|
||
samples = samples.reshape(1, -1)
|
||
elif samples.ndim == 2 and samples.shape[0] > samples.shape[1]:
|
||
samples = np.ascontiguousarray(samples.T)
|
||
num_channels = samples.shape[0]
|
||
layout = "stereo" if num_channels >= 2 else "mono"
|
||
a_stream = cast(av.AudioStream, container.add_stream(audio_codec, rate=audio_sample_rate))
|
||
a_stream.layout = layout
|
||
|
||
for frame_data in video_frames:
|
||
frame = av.VideoFrame.from_ndarray(frame_data, format="rgb24")
|
||
for packet in v_stream.encode(frame):
|
||
container.mux(packet)
|
||
for packet in v_stream.encode():
|
||
container.mux(packet)
|
||
|
||
if a_stream is not None and audio_waveform is not None:
|
||
if samples is None or layout is None:
|
||
raise ValueError("Audio samples were not prepared for muxing.")
|
||
audio_frame = av.AudioFrame.from_ndarray(samples, format="fltp", layout=layout)
|
||
audio_frame.sample_rate = audio_sample_rate
|
||
# AAC has a one-frame encoder delay. Mark the input waveform as
|
||
# starting at t=0 so the MP4 muxer writes the corresponding negative
|
||
# priming timestamp instead of exposing the delay as leading silence.
|
||
audio_frame.pts = 0
|
||
audio_frame.time_base = Fraction(1, audio_sample_rate)
|
||
for packet in a_stream.encode(audio_frame):
|
||
container.mux(packet)
|
||
for packet in a_stream.encode():
|
||
container.mux(packet)
|
||
|
||
container.close()
|
||
return buf.getvalue()
|
||
|
||
|
||
def mux_av_video_audio_bytes(
|
||
video_frames: Iterable[av.VideoFrame],
|
||
width: int,
|
||
height: int,
|
||
audio_waveform: np.ndarray | None = None,
|
||
*,
|
||
fps: float = 25.0,
|
||
audio_sample_rate: int | None = None,
|
||
video_codec: str = "h264",
|
||
audio_codec: str = "aac",
|
||
crf: str = "18",
|
||
video_codec_options: dict[str, str] | None = None,
|
||
) -> bytes:
|
||
"""Mux preconstructed video frames and optional audio into MP4 bytes."""
|
||
buf = io.BytesIO()
|
||
with cast(Any, av.open(buf, mode="w", format="mp4")) as container:
|
||
v_stream = cast(
|
||
av.VideoStream,
|
||
container.add_stream(video_codec, rate=Fraction(fps).limit_denominator(10000)),
|
||
)
|
||
v_stream.width = width
|
||
v_stream.height = height
|
||
v_stream.pix_fmt = "yuv420p"
|
||
|
||
options: dict[str, object] = {"crf": str(crf)}
|
||
if video_codec_options:
|
||
options.update(video_codec_options)
|
||
v_stream.options = options
|
||
|
||
a_stream: av.AudioStream | None = None
|
||
samples: np.ndarray | None = None
|
||
layout: str | None = None
|
||
if audio_waveform is not None:
|
||
effective_audio_sample_rate = 44100 if audio_sample_rate is None else audio_sample_rate
|
||
samples = audio_waveform.astype(np.float32)
|
||
if samples.ndim == 1:
|
||
samples = samples.reshape(1, -1)
|
||
elif samples.ndim == 2 and samples.shape[0] > samples.shape[1]:
|
||
samples = np.ascontiguousarray(samples.T)
|
||
num_channels = samples.shape[0]
|
||
layout = "stereo" if num_channels >= 2 else "mono"
|
||
a_stream = cast(av.AudioStream, container.add_stream(audio_codec, rate=effective_audio_sample_rate))
|
||
a_stream.layout = layout
|
||
|
||
for frame in video_frames:
|
||
for packet in v_stream.encode(frame):
|
||
container.mux(packet)
|
||
for packet in v_stream.encode():
|
||
container.mux(packet)
|
||
|
||
if a_stream is not None and audio_waveform is not None:
|
||
if samples is None or layout is None:
|
||
raise ValueError("Audio samples were not prepared for muxing.")
|
||
audio_frame = av.AudioFrame.from_ndarray(samples, format="fltp", layout=layout)
|
||
audio_frame.sample_rate = effective_audio_sample_rate
|
||
# AAC has a one-frame encoder delay. Mark the input waveform as
|
||
# starting at t=0 so the MP4 muxer writes the corresponding negative
|
||
# priming timestamp instead of exposing the delay as leading silence.
|
||
audio_frame.pts = 0
|
||
audio_frame.time_base = Fraction(1, effective_audio_sample_rate)
|
||
for packet in a_stream.encode(audio_frame):
|
||
container.mux(packet)
|
||
for packet in a_stream.encode():
|
||
container.mux(packet)
|
||
|
||
return buf.getvalue()
|