2ab5d1701a
Signed-off-by: specture724 <specture724@gmail.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
88 lines
3.9 KiB
Python
88 lines
3.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Worker-side MP4 encoding for Wan VAE temporal chunks."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from vllm_omni.diffusion.utils.media_utils import normalize_preencode_batch_frames, normalize_video_codec_options
|
|
|
|
# Matches the fallback the video serving layer applies when a request omits fps,
|
|
# so the bytes the worker encodes carry the rate the response advertises.
|
|
WAN_DEFAULT_OUTPUT_FPS = 24
|
|
|
|
# Keep worker-side encoding consistent with the one-shot serving encoder when
|
|
# a request does not choose codec options explicitly.
|
|
WAN_DEFAULT_VIDEO_CODEC_OPTIONS = {"preset": "ultrafast", "threads": "0"}
|
|
|
|
# Wan publishes one latent frame group per callback, which is far finer than one
|
|
# host transfer is worth. 17 frames matches MiniMax-H3's native clip size.
|
|
WAN_DEFAULT_BATCH_FRAMES = 17
|
|
|
|
|
|
def wan_preencoded_mp4_payload(video: Any) -> dict[str, Any] | None:
|
|
"""Return a post-process payload when the worker already encoded the MP4.
|
|
|
|
Returns ``None`` for the ordinary frame-tensor output so callers fall
|
|
through to their existing post-processing.
|
|
|
|
No fps is reported here: :func:`resolve_wan_output_fps` mirrors the fallback
|
|
the serving layer applies, so the rate the response advertises already
|
|
matches the rate these bytes were encoded at.
|
|
"""
|
|
if isinstance(video, (bytes, bytearray, memoryview)):
|
|
video = [video]
|
|
if not isinstance(video, list) or not all(isinstance(item, (bytes, bytearray, memoryview)) for item in video):
|
|
return None
|
|
return {"payload": {"video": [bytes(item) for item in video]}, "metadata": {}}
|
|
|
|
|
|
def resolve_wan_preencode_mp4(sampling_params: Any, *, output_type: str) -> bool:
|
|
"""Return whether this request wants the worker to pre-encode the MP4.
|
|
|
|
Reject rather than silently downgrade the combinations the pre-encoded path
|
|
cannot serve: it emits MP4 bytes, so nothing downstream can interpolate
|
|
frames or hand back an array.
|
|
"""
|
|
extra_args = getattr(sampling_params, "extra_args", None) or {}
|
|
if not extra_args.get("preencode_mp4", False):
|
|
return False
|
|
if getattr(sampling_params, "enable_frame_interpolation", False):
|
|
raise ValueError(
|
|
"preencode_mp4 cannot be combined with enable_frame_interpolation: interpolation "
|
|
"needs the decoded frames that the pre-encoded path no longer materializes"
|
|
)
|
|
if output_type != "np":
|
|
raise ValueError(f"preencode_mp4 returns MP4 bytes and cannot serve output_type={output_type!r}")
|
|
return True
|
|
|
|
|
|
def resolve_wan_preencode_batch_frames(sampling_params: Any, *, default: int = WAN_DEFAULT_BATCH_FRAMES) -> int:
|
|
"""Read the MP4 batching threshold before starting expensive generation."""
|
|
extra_args = getattr(sampling_params, "extra_args", None) or {}
|
|
return normalize_preencode_batch_frames(extra_args.get("preencode_batch_frames", default))
|
|
|
|
|
|
def resolve_wan_video_codec_options(sampling_params: Any) -> dict[str, str] | None:
|
|
"""Read the request's encoder options, preserving serving defaults.
|
|
|
|
An absent key inherits the one-shot serving encoder's defaults. Explicit
|
|
``None`` and an empty mapping remain distinct caller choices.
|
|
"""
|
|
extra_args = getattr(sampling_params, "extra_args", None) or {}
|
|
return normalize_video_codec_options(extra_args.get("video_codec_options", WAN_DEFAULT_VIDEO_CODEC_OPTIONS))
|
|
|
|
|
|
def resolve_wan_output_fps(sampling_params: Any) -> int:
|
|
"""Resolve the fps the worker encodes with.
|
|
|
|
The serving layer resolves output fps from the request's ``fps`` field, not
|
|
from ``frame_rate``, so pre-encoding has to read the same field or the bytes
|
|
would carry a rate the response does not advertise.
|
|
"""
|
|
fps = getattr(sampling_params, "fps", None)
|
|
if isinstance(fps, list):
|
|
fps = fps[0] if fps else None
|
|
return int(fps) if fps else WAN_DEFAULT_OUTPUT_FPS
|