class VideoMediaIO(MediaIO[tuple[DecodedFrames, dict[str, Any]]]):
"""Configuration values can be user-provided either by --media-io-kwargs or
by the runtime API field "media_io_kwargs". Ensure proper validation and
error handling.
Use --media-io-kwargs '{"video": {"backend": "torchcodec",
"device": "cuda"}}' to decode with NVDEC and keep frames on the GPU,
which pairs with a device-side multi-modal processor
(--mm-processor-device, EPD encode-only instances).
"""
@classmethod
def merge_kwargs(
cls,
default_kwargs: dict[str, Any] | None,
runtime_kwargs: dict[str, Any] | None,
) -> dict[str, Any]:
if runtime_kwargs:
# Decoder GPU memory is reserved from the startup value.
runtime_kwargs = dict(runtime_kwargs)
runtime_kwargs.pop("hw_decoders", None)
runtime_kwargs.pop("pool_size", None)
# The decode device determines which device the decoded frames
# (and thus the processor pipeline) live on; startup-only.
runtime_kwargs.pop("device", None)
# Block request-level selection of GPU video backends that
# were not configured (and VRAM-reserved) at startup.
for key in ("video_backend", "backend"):
requested = runtime_kwargs.get(key)
if requested and VIDEO_LOADER_REGISTRY.backend_requires_gpu(requested):
static_val = (default_kwargs or {}).get(key)
if static_val != requested:
logger.warning_once(
"Stripping request-level %s=%r: GPU video "
"backend not configured at startup.",
key,
requested,
)
runtime_kwargs = {
k: v for k, v in runtime_kwargs.items() if k != key
}
merged = super().merge_kwargs(default_kwargs, runtime_kwargs)
# fps and num_frames interact with each other, so if either is
# overridden at request time, wipe the other from defaults to
# avoid unintuitive cross-field interactions.
if runtime_kwargs:
if "num_frames" in runtime_kwargs and "fps" not in runtime_kwargs:
merged.pop("fps", None)
elif "fps" in runtime_kwargs and "num_frames" not in runtime_kwargs:
merged.pop("num_frames", None)
requested_backend = runtime_kwargs.get("backend")
static_backend = (default_kwargs or {}).get("backend")
if requested_backend and requested_backend != static_backend:
for name in incompatible_backend_options(
static_backend, requested_backend
):
merged.pop(name, None)
return merged
def __init__(
self,
image_io: ImageMediaIO,
num_frames: int = 32,
**kwargs,
) -> None:
super().__init__()
self.image_io = image_io
self.num_frames = num_frames
# `kwargs` contains custom arguments from
# --media-io-kwargs for this modality, merged with
# per-request runtime media_io_kwargs via merge_kwargs().
# They can be passed to the underlying
# media loaders (e.g. custom implementations)
# for flexible control.
# Allow per-request override of video backend via kwargs.
# This enables users to specify a different backend than the
# global VLLM_VIDEO_LOADER_BACKEND env var, e.g.:
# --media-io-kwargs '{"video": {"video_backend": "torchcodec"}}'
video_loader_backend = (
kwargs.pop("video_backend", None) or envs.VLLM_VIDEO_LOADER_BACKEND
)
self.kwargs = kwargs
self.video_loader_backend = video_loader_backend
self.video_loader = VIDEO_LOADER_REGISTRY.load(video_loader_backend)
def get_decode_spec(self) -> DecodeSpec:
return DecodeSpec(
{
**self.kwargs,
"num_frames": self.num_frames,
"video_backend": self.video_loader_backend,
}
)
def load_bytes(self, data: bytes) -> tuple[DecodedFrames, dict[str, Any]]:
return self.video_loader.load_bytes(
data, num_frames=self.num_frames, **self.kwargs
)
def _prepare_jpeg_sequence(
self, data: str
) -> tuple[Callable[[], npt.NDArray], dict[str, Any]]:
"""Split a base64 jpeg frame sequence and validate its arguments.
Splitting and argument validation (num_frames/frames_indices/
total_num_frames/duration) run eagerly; the per-frame decode and
np.stack are deferred into the returned callable.
"""
if self.num_frames > 0:
frame_parts = data.split(",", self.num_frames)[: self.num_frames]
elif self.num_frames == 0:
raise ValueError("num_frames must be greater than 0 or -1")
else:
frame_parts = data.split(",")
total = len(frame_parts)
fps = float(self.kwargs.get("fps", 1))
# validate and extract frames_indices
frames_indices = self.kwargs.get("frames_indices")
if frames_indices is not None:
if not (
isinstance(frames_indices, list)
and all(isinstance(i, int) for i in frames_indices)
):
raise ValueError("frames_indices must be a list of integers")
if len(frames_indices) != total:
raise ValueError(
f"frames_indices length ({len(frames_indices)}) must "
f"match number of frames sent ({total})"
)
else:
frames_indices = list(range(total))
# validate and extract total_num_frames
total_num_frames = self.kwargs.get("total_num_frames", total)
if not isinstance(total_num_frames, int) or total_num_frames < 1:
raise ValueError("total_num_frames must be a positive integer")
if total_num_frames < total:
raise ValueError(
f"total_num_frames ({total_num_frames}) must be >= "
f"number of frames sent ({total})"
)
# validate and extract duration
duration = self.kwargs.get("duration")
if duration is not None:
if not isinstance(duration, (int, float)) or duration < 0:
raise ValueError("duration must be a non-negative number")
else:
duration = total_num_frames / fps if fps > 0 else 0.0
metadata = {
"total_num_frames": total_num_frames,
"fps": fps,
"duration": duration,
"video_backend": "jpeg_sequence",
"frames_indices": frames_indices,
"do_sample_frames": self.kwargs.get("do_sample_frames", False),
}
def decode_frames() -> npt.NDArray:
load_frame = partial(
self.image_io.load_base64,
"image/jpeg",
)
return np.stack(
[np.asarray(load_frame(frame_data)) for frame_data in frame_parts]
)
return decode_frames, metadata
def load_base64(
self, media_type: str, data: str
) -> tuple[DecodedFrames, dict[str, Any]]:
if media_type.lower() == "video/jpeg":
decode_frames, metadata = self._prepare_jpeg_sequence(data)
return decode_frames(), metadata
return self.load_bytes(pybase64.b64decode(data, validate=True))
def load_base64_ref(
self, media_type: str, data: str
) -> MediaRef[tuple[DecodedFrames, dict[str, Any]]]:
if media_type.lower() == "video/jpeg":
decode_frames, metadata = self._prepare_jpeg_sequence(data)
spec = self.get_decode_spec().extend(
video_backend="jpeg_sequence",
image_decode=dict(self.image_io.get_decode_spec().settings),
)
return MediaRef(
lambda: (decode_frames(), metadata),
data.encode(),
spec,
)
return super().load_base64_ref(media_type, data)
def load_file(self, filepath: Path) -> tuple[DecodedFrames, dict[str, Any]]:
return self.load_bytes(filepath.read_bytes())
def encode_base64(
self,
media: npt.NDArray,
*,
video_format: str = "JPEG",
) -> str:
video = media
if video_format == "JPEG":
encode_frame = partial(
self.image_io.encode_base64,
image_format=video_format,
)
return ",".join(encode_frame(Image.fromarray(frame)) for frame in video)
msg = "Only JPEG format is supported for now."
raise NotImplementedError(msg)