[Experimental][Perf] Eliminate redundant host copies from generated-video transport (#7355)
Signed-off-by: mglyn <1203789601@qq.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
1 parent
c27a52a359
commit
3a82c85880
10 files changed
+339
-34
No files matched your search
@@ -70,3 +70,42 @@ as safety checks or audio/video packaging.
|
||||
|
||||
See [RFC #6541](https://github.com/vllm-project/vllm-omni/issues/6541) for the
|
||||
contract, lifecycle, batching rules, and migration plan.
|
||||
|
||||
## Experimental registered-SHM video transport
|
||||
|
||||
`video_output_transport.enable_registered_shm` is a separate, default-off option
|
||||
for CUDA request-mode workers. It copies eligible video tensors directly into
|
||||
CUDA-registered shared memory and returns borrowed CPU views, removing the
|
||||
intermediate pinned-to-SHM and SHM-to-consumer host copies. The D2H transfer still
|
||||
occurs.
|
||||
|
||||
`video_output_transport.enable_borrowed_frames` independently enables borrowed
|
||||
RGB AVFrames in the Videos API MP4 response encoder, removing the third host
|
||||
copy into PyAV frame buffers. It also defaults to false. Eligible uint8 RGB
|
||||
frames are wrapped with `VideoFrame.from_numpy_buffer`; other dtypes/layouts use
|
||||
the existing encoder path. Color conversion, compression and muxing still occur.
|
||||
|
||||
Enable it through the same CLI, Python or deploy-stage configuration:
|
||||
|
||||
```bash
|
||||
--video-output-transport '{"enable_registered_shm": true, "enable_borrowed_frames": true}'
|
||||
```
|
||||
|
||||
Registered SHM applies only to prepared typed video media (currently WAN2.2) and
|
||||
explicitly marked legacy video outputs (currently MiniMax-H3). H3 marks the video
|
||||
entry of its `(video, audio)` output; audio, trajectory data, images and unmarked
|
||||
legacy outputs keep their existing transport. Device-side postprocessing is
|
||||
independent and does not need to be enabled for registered SHM.
|
||||
Pre-encoded MP4 outputs bypass raw-frame transport and API-side encoding.
|
||||
|
||||
The existing SHM routing threshold is unchanged: the tensor or its backing
|
||||
storage must exceed 1,000,000 bytes. This is not a measured performance crossover.
|
||||
SHM is allocated and registered on every transfer, without pooling; small outputs
|
||||
can regress and registration latency can vary. Benchmark representative output
|
||||
sizes before opting in. CPU outputs, non-CUDA backends, and synchronous/step-mode
|
||||
transfers keep their existing path.
|
||||
|
||||
Borrowed views keep the unlinked mapping alive until their last user releases
|
||||
it. CUDA registration/copy failures fail the request; if CUDA cleanup cannot be
|
||||
confirmed, the worker retains the mapping for safety and must be restarted before
|
||||
further registered-SHM transfers.
|
||||
@@ -151,6 +151,48 @@ def test_raw_and_base64_encoders_receive_persistent_converter(mocker: MockerFixt
|
||||
handler.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("typed_stage", [False, True], ids=["legacy", "typed"])
|
||||
def test_remote_diffusion_stage_enables_borrowed_frames_for_both_video_responses(
|
||||
mocker: MockerFixture, typed_stage: bool
|
||||
):
|
||||
from vllm_omni.config.config_factory import StageConfigFactory
|
||||
|
||||
options = {
|
||||
"model_class_name": "MiniMaxH3Pipeline",
|
||||
"video_output_transport": {"enable_borrowed_frames": True},
|
||||
}
|
||||
stage = (
|
||||
StageConfigFactory.create_typed_default_diffusion("test-model", options).stage_configs[0]
|
||||
if typed_stage
|
||||
else StageConfigFactory.create_default_diffusion(options)[0]
|
||||
)
|
||||
# Like StageDiffusionClient, this frontend engine exposes metadata but no
|
||||
# worker od_config. The effective setting is available in stage_configs.
|
||||
engine = FakeAsyncOmni()
|
||||
engine.stage_configs = [stage]
|
||||
handler = OmniOpenAIServingVideo.for_diffusion(engine, model_name="test-model", stage_configs=engine.stage_configs)
|
||||
raw_encoder = mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
|
||||
return_value=b"encoded-video",
|
||||
)
|
||||
base64_encoder = mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
|
||||
return_value="encoded-video",
|
||||
)
|
||||
try:
|
||||
|
||||
async def generate_both():
|
||||
request = VideoGenerationRequest(prompt="test prompt")
|
||||
await handler.generate_video_bytes(request, "raw-request")
|
||||
await handler.generate_videos(request, "base64-request")
|
||||
|
||||
asyncio.run(generate_both())
|
||||
assert raw_encoder.call_args.kwargs["enable_borrowed_frames"] is True
|
||||
assert base64_encoder.call_args.kwargs["enable_borrowed_frames"] is True
|
||||
finally:
|
||||
handler.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_frames", [0, -1, True, 1.5, "17", None])
|
||||
def test_preencode_rejects_invalid_batch_frames_before_generation(batch_frames):
|
||||
engine = FakeAsyncOmni()
|
||||
|
||||
@@ -807,10 +807,16 @@ def uses_diffusers_adapter(od_config: object) -> bool:
|
||||
@dataclass
|
||||
class VideoOutputTransportConfig:
|
||||
enable_device_postprocess: bool = False
|
||||
enable_registered_shm: bool = False
|
||||
enable_borrowed_frames: bool = False
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if not isinstance(self.enable_device_postprocess, bool):
|
||||
raise TypeError("enable_device_postprocess must be a bool")
|
||||
if not isinstance(self.enable_registered_shm, bool):
|
||||
raise TypeError("enable_registered_shm must be a bool")
|
||||
if not isinstance(self.enable_borrowed_frames, bool):
|
||||
raise TypeError("enable_borrowed_frames must be a bool")
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -1826,11 +1832,16 @@ class DiffusionOutput:
|
||||
# Internal control-plane event emitted on first scheduler admission.
|
||||
request_started: bool = False
|
||||
|
||||
# Typed video-media contract. Declared last so the pre-existing positional
|
||||
# Typed video-media contract. Appended so the pre-existing positional
|
||||
# constructor order (output, trajectory_timesteps, ...) that out-of-tree
|
||||
# pipelines rely on is preserved. Mutually exclusive with ``output``.
|
||||
media: DiffusionMediaOutput | None = None
|
||||
|
||||
# Compatibility adapter for joint (video, audio) outputs that have not
|
||||
# migrated to typed media. Only this tuple entry is eligible for video
|
||||
# transport optimizations; unmarked legacy outputs retain their old path.
|
||||
video_output_index: int | None = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.media is not None and not isinstance(self.media, DiffusionMediaOutput):
|
||||
raise TypeError(f"media must be DiffusionMediaOutput, got {type(self.media).__name__}")
|
||||
|
||||
+168
-26
@@ -12,10 +12,12 @@ serialised through the queue.
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterator
|
||||
from typing import Any
|
||||
from contextlib import suppress
|
||||
from typing import Any, TypeGuard
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput
|
||||
from vllm_omni.diffusion.media import (
|
||||
@@ -33,49 +35,144 @@ from vllm_omni.diffusion.media import (
|
||||
_SHM_TENSOR_THRESHOLD = 1_000_000 # 1 MB
|
||||
DIFFUSION_RPC_RESULT_ENVELOPE = "diffusion_rpc_result"
|
||||
_DIFFUSION_MEDIA_WIRE_TYPE = "diffusion_media_v1"
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# An uncertain CUDA cleanup must not unmap memory still used by the device.
|
||||
# Retain failed transfers until worker exit and reject further registrations.
|
||||
_failed_shm_copies: list[tuple[Any, torch.Tensor, torch.Tensor, torch.Stream]] = []
|
||||
|
||||
# Sentinel so compute-then-assign packing can tell "field not packed" apart from
|
||||
# "field packed to None".
|
||||
_UNSET = object()
|
||||
|
||||
|
||||
def _array_to_shm(array: np.ndarray) -> dict[str, Any]:
|
||||
"""Copy a contiguous NumPy-compatible array into shared memory."""
|
||||
class _SharedMemoryLease:
|
||||
"""Keep an unlinked shared-memory mapping alive while array views use it."""
|
||||
|
||||
def __init__(self, shm: Any, array: np.ndarray) -> None:
|
||||
self.shm = shm
|
||||
self.__array_interface__ = array.__array_interface__
|
||||
|
||||
def __del__(self) -> None:
|
||||
with suppress(BufferError):
|
||||
self.shm.close()
|
||||
|
||||
|
||||
def _new_shm_array(shape: tuple[int, ...], dtype: np.dtype) -> tuple[Any, np.ndarray, dict[str, Any]]:
|
||||
from multiprocessing import shared_memory
|
||||
|
||||
nbytes = int(np.prod(shape, dtype=np.int64)) * dtype.itemsize
|
||||
shm = shared_memory.SharedMemory(create=True, size=nbytes)
|
||||
array = np.ndarray(shape, dtype=dtype, buffer=shm.buf[:nbytes])
|
||||
handle = {
|
||||
"name": shm.name,
|
||||
"shape": list(shape),
|
||||
"numpy_dtype": str(dtype),
|
||||
"nbytes": nbytes,
|
||||
}
|
||||
return shm, array, handle
|
||||
|
||||
|
||||
def _array_to_shm(array: np.ndarray) -> dict[str, Any]:
|
||||
"""Copy a contiguous NumPy-compatible array into shared memory."""
|
||||
if array.dtype.hasobject:
|
||||
raise TypeError("NumPy object arrays cannot be transferred through raw shared memory")
|
||||
|
||||
array = np.ascontiguousarray(array)
|
||||
nbytes = array.nbytes
|
||||
shm = shared_memory.SharedMemory(create=True, size=nbytes)
|
||||
shm_array = np.ndarray(array.shape, dtype=array.dtype, buffer=shm.buf[:nbytes])
|
||||
shm, shm_array, handle = _new_shm_array(array.shape, array.dtype)
|
||||
np.copyto(shm_array, array)
|
||||
handle = {
|
||||
"name": shm.name,
|
||||
"shape": list(array.shape),
|
||||
"numpy_dtype": str(array.dtype),
|
||||
"nbytes": nbytes,
|
||||
}
|
||||
del shm_array
|
||||
shm.close()
|
||||
return handle
|
||||
|
||||
|
||||
def _array_from_shm(handle: dict[str, Any]) -> np.ndarray:
|
||||
"""Copy an array from shared memory, then close and unlink its segment."""
|
||||
"""Materialize or borrow an array from shared memory and unlink its name."""
|
||||
from multiprocessing import shared_memory
|
||||
|
||||
shm = shared_memory.SharedMemory(name=handle["name"])
|
||||
array = np.ndarray(
|
||||
handle["shape"],
|
||||
dtype=np.dtype(handle["numpy_dtype"]),
|
||||
buffer=shm.buf[: handle["nbytes"]],
|
||||
)
|
||||
if handle.get("borrow_on_unpack"):
|
||||
# A base owner survives np.asarray(), slicing, and torch.from_numpy().
|
||||
# An ndarray subclass attribute does not survive np.asarray().
|
||||
borrowed = np.asarray(_SharedMemoryLease(shm, array))
|
||||
shm.unlink()
|
||||
return borrowed
|
||||
try:
|
||||
array = np.ndarray(
|
||||
handle["shape"],
|
||||
dtype=np.dtype(handle["numpy_dtype"]),
|
||||
buffer=shm.buf[: handle["nbytes"]],
|
||||
).copy()
|
||||
return array.copy()
|
||||
finally:
|
||||
del array
|
||||
shm.close()
|
||||
shm.unlink()
|
||||
return array
|
||||
|
||||
|
||||
def _cuda_registered_shm_copy(tensor: torch.Tensor, d2h_stream: torch.Stream) -> dict[str, Any]:
|
||||
"""Copy a CUDA tensor directly into a CUDA-registered POSIX SHM segment."""
|
||||
if _failed_shm_copies:
|
||||
raise RuntimeError("A previous registered SHM transfer failed CUDA cleanup; restart the worker")
|
||||
original_dtype = tensor.dtype
|
||||
source = tensor.detach()
|
||||
# NumPy has no BF16 dtype. Transfer the original bits, without allocating
|
||||
# an FP32 accelerator tensor or converting the borrowed CPU result.
|
||||
if original_dtype == torch.bfloat16:
|
||||
numpy_dtype = np.dtype(np.uint16)
|
||||
else:
|
||||
numpy_dtype = torch.empty((), dtype=source.dtype).numpy().dtype
|
||||
cudart = torch.cuda.cudart()
|
||||
shm, shm_array, handle = _new_shm_array(tuple(source.shape), numpy_dtype)
|
||||
pointer = shm_array.ctypes.data
|
||||
host_tensor = None
|
||||
registered = False
|
||||
succeeded = False
|
||||
try:
|
||||
host_tensor = torch.from_numpy(shm_array).view(original_dtype)
|
||||
error = cudart.cudaHostRegister(pointer, handle["nbytes"], 0)
|
||||
if int(error) != 0:
|
||||
raise RuntimeError(f"cudaHostRegister failed: {cudart.cudaGetErrorString(error)}")
|
||||
registered = True
|
||||
old_stream = torch.accelerator.current_stream()
|
||||
torch.accelerator.set_stream(d2h_stream)
|
||||
try:
|
||||
host_tensor.copy_(source, non_blocking=True)
|
||||
finally:
|
||||
torch.accelerator.set_stream(old_stream)
|
||||
d2h_stream.synchronize()
|
||||
error = cudart.cudaHostUnregister(pointer)
|
||||
if int(error) != 0:
|
||||
raise RuntimeError(f"cudaHostUnregister failed: {cudart.cudaGetErrorString(error)}")
|
||||
registered = False
|
||||
handle["borrow_on_unpack"] = True
|
||||
handle.update(
|
||||
{
|
||||
"__tensor_shm__": True,
|
||||
"torch_dtype": str(original_dtype),
|
||||
"torch_dtype_view": original_dtype == torch.bfloat16,
|
||||
}
|
||||
)
|
||||
succeeded = True
|
||||
return handle
|
||||
finally:
|
||||
if registered:
|
||||
try:
|
||||
# A copy may have been enqueued before a later operation failed.
|
||||
d2h_stream.synchronize()
|
||||
error = cudart.cudaHostUnregister(pointer)
|
||||
if int(error) != 0:
|
||||
raise RuntimeError(f"cudaHostUnregister failed: {cudart.cudaGetErrorString(error)}")
|
||||
registered = False
|
||||
except Exception:
|
||||
assert host_tensor is not None
|
||||
_failed_shm_copies.append((shm, host_tensor, source, d2h_stream))
|
||||
logger.exception("Retaining registered SHM after CUDA cleanup failure; worker restart required")
|
||||
del host_tensor, shm_array
|
||||
if not registered:
|
||||
shm.close()
|
||||
if not succeeded:
|
||||
shm.unlink()
|
||||
|
||||
|
||||
def _unlink_shm_handle(handle: dict[str, Any]) -> None:
|
||||
@@ -111,6 +208,8 @@ def _tensor_to_shm(
|
||||
tensor: torch.Tensor,
|
||||
d2h_stream: torch.Stream | None = None,
|
||||
created: list[dict[str, Any]] | None = None,
|
||||
*,
|
||||
enable_registered_shm: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
"""Copy a tensor into POSIX shared memory and return a metadata handle.
|
||||
|
||||
@@ -123,6 +222,17 @@ def _tensor_to_shm(
|
||||
packed.
|
||||
"""
|
||||
original_dtype = tensor.dtype
|
||||
if (
|
||||
d2h_stream is not None
|
||||
and enable_registered_shm
|
||||
and tensor.device.type == "cuda"
|
||||
# ROCm also uses the "cuda" device type, but not this CUDA runtime API.
|
||||
and torch.version.cuda is not None
|
||||
):
|
||||
handle = _cuda_registered_shm_copy(tensor, d2h_stream)
|
||||
if created is not None:
|
||||
created.append(handle)
|
||||
return handle
|
||||
if d2h_stream is not None:
|
||||
# Non-blocking D2H: copy on side stream to pinned CPU memory.
|
||||
old_stream = torch.accelerator.current_stream()
|
||||
@@ -165,7 +275,7 @@ def _tensor_from_shm(handle: dict[str, Any]) -> torch.Tensor:
|
||||
if torch_dtype_str:
|
||||
original_dtype = getattr(torch, torch_dtype_str.replace("torch.", ""), None)
|
||||
if original_dtype is not None and tensor.dtype != original_dtype:
|
||||
tensor = tensor.to(original_dtype)
|
||||
tensor = tensor.view(original_dtype) if handle.get("torch_dtype_view") else tensor.to(original_dtype)
|
||||
return tensor
|
||||
|
||||
|
||||
@@ -173,6 +283,8 @@ def _pack_tensor_if_large(
|
||||
val: torch.Tensor,
|
||||
d2h_stream: torch.Stream | None = None,
|
||||
created: list[dict[str, Any]] | None = None,
|
||||
*,
|
||||
enable_registered_shm: bool = False,
|
||||
) -> torch.Tensor | dict:
|
||||
"""Replace a tensor with an SHM handle if it exceeds the threshold.
|
||||
|
||||
@@ -187,7 +299,7 @@ def _pack_tensor_if_large(
|
||||
except Exception:
|
||||
storage_bytes = view_bytes
|
||||
if max(view_bytes, storage_bytes) > _SHM_TENSOR_THRESHOLD:
|
||||
return _tensor_to_shm(val, d2h_stream=d2h_stream, created=created)
|
||||
return _tensor_to_shm(val, d2h_stream=d2h_stream, created=created, enable_registered_shm=enable_registered_shm)
|
||||
return val
|
||||
|
||||
|
||||
@@ -255,12 +367,15 @@ def _pack_diffusion_media(
|
||||
*,
|
||||
d2h_stream: torch.Stream | None,
|
||||
created: list[dict[str, Any]] | None = None,
|
||||
enable_registered_shm: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
media.validate()
|
||||
if not media.prepared_for_transport:
|
||||
raise ValueError("Diffusion media must be prepared before IPC packing")
|
||||
video = media.video
|
||||
packed_tensor = _pack_tensor_if_large(video.tensor, d2h_stream=d2h_stream, created=created)
|
||||
packed_tensor = _pack_tensor_if_large(
|
||||
video.tensor, d2h_stream=d2h_stream, created=created, enable_registered_shm=enable_registered_shm
|
||||
)
|
||||
if isinstance(packed_tensor, torch.Tensor) and packed_tensor.device.type != "cpu":
|
||||
packed_tensor = packed_tensor.detach().cpu().contiguous()
|
||||
return {
|
||||
@@ -319,6 +434,8 @@ def _pack_diffusion_fields(
|
||||
d2h_stream: torch.Stream | None = None,
|
||||
created: list[dict[str, Any]] | None = None,
|
||||
pending_updates: list[tuple[DiffusionOutput, str, object]] | None = None,
|
||||
*,
|
||||
enable_registered_shm: bool = False,
|
||||
) -> DiffusionOutput:
|
||||
if output.media is not None and output.output is not None:
|
||||
raise ValueError("DiffusionOutput cannot contain both media and legacy output")
|
||||
@@ -329,11 +446,26 @@ def _pack_diffusion_fields(
|
||||
# (never enqueue a half-packed object); the caller unlinks any SHM segments
|
||||
# accumulated in ``created`` on failure.
|
||||
packed_output = _UNSET
|
||||
if output.output is not None:
|
||||
if enable_registered_shm and output.video_output_index is not None:
|
||||
index = output.video_output_index
|
||||
values = output.output
|
||||
if not isinstance(values, tuple) or type(index) is not int or not 0 <= index < len(values):
|
||||
raise ValueError("video_output_index must identify a video tensor in the legacy output tuple")
|
||||
if not isinstance(values[index], torch.Tensor):
|
||||
raise TypeError("The legacy video output must be a tensor")
|
||||
packed_output = tuple(
|
||||
_pack_tensor_if_large(value, d2h_stream=d2h_stream, created=created, enable_registered_shm=True)
|
||||
if position == index
|
||||
else _pack_value_if_large(value, d2h_stream=d2h_stream, created=created)
|
||||
for position, value in enumerate(values)
|
||||
)
|
||||
elif output.output is not None:
|
||||
packed_output = _pack_value_if_large(output.output, d2h_stream=d2h_stream, created=created)
|
||||
packed_media = _UNSET
|
||||
if output.media is not None:
|
||||
packed_media = _pack_diffusion_media(output.media, d2h_stream=d2h_stream, created=created)
|
||||
packed_media = _pack_diffusion_media(
|
||||
output.media, d2h_stream=d2h_stream, created=created, enable_registered_shm=enable_registered_shm
|
||||
)
|
||||
packed_latents = _UNSET
|
||||
if isinstance(output.trajectory_latents, torch.Tensor):
|
||||
packed_latents = _pack_tensor_if_large(output.trajectory_latents, d2h_stream=d2h_stream, created=created)
|
||||
@@ -363,13 +495,15 @@ def _pack_diffusion_fields(
|
||||
return output
|
||||
|
||||
|
||||
def _is_rpc_result_envelope(output: object) -> bool:
|
||||
def _is_rpc_result_envelope(output: object) -> TypeGuard[dict[str, Any]]:
|
||||
return isinstance(output, dict) and output.get("type") == DIFFUSION_RPC_RESULT_ENVELOPE
|
||||
|
||||
|
||||
def pack_diffusion_output_shm(
|
||||
output: object,
|
||||
d2h_stream: torch.Stream | None = None,
|
||||
*,
|
||||
enable_registered_shm: bool = False,
|
||||
) -> object:
|
||||
"""Replace large tensors in diffusion worker outputs with SHM handles.
|
||||
|
||||
@@ -382,6 +516,10 @@ def pack_diffusion_output_shm(
|
||||
If *d2h_stream* is provided, D2H copies use that stream (non-blocking on
|
||||
the default stream). The caller must synchronize *d2h_stream* afterward.
|
||||
|
||||
``enable_registered_shm`` opts explicitly identified video tensors into
|
||||
direct CUDA-to-SHM transfer and borrowed CPU views. Other payloads keep
|
||||
their existing transport. Registration is per transfer, without pooling.
|
||||
|
||||
Packing is failure-atomic: every SHM segment created during the call is
|
||||
tracked and unlinked if any part of the payload fails to pack, and field
|
||||
mutations are committed only after the complete payload succeeds.
|
||||
@@ -394,6 +532,7 @@ def pack_diffusion_output_shm(
|
||||
d2h_stream=d2h_stream,
|
||||
created=created,
|
||||
pending_updates=pending_updates,
|
||||
enable_registered_shm=enable_registered_shm,
|
||||
)
|
||||
except BaseException:
|
||||
_unlink_shm_handles(created)
|
||||
@@ -444,13 +583,16 @@ def _pack_output_shm(
|
||||
d2h_stream: torch.Stream | None,
|
||||
created: list[dict[str, Any]],
|
||||
pending_updates: list[tuple[DiffusionOutput, str, object]],
|
||||
enable_registered_shm: bool = False,
|
||||
) -> object:
|
||||
transport_options = {"enable_registered_shm": True} if enable_registered_shm else {}
|
||||
for item in _iter_diffusion_outputs(output):
|
||||
_pack_diffusion_fields(
|
||||
item,
|
||||
d2h_stream=d2h_stream,
|
||||
created=created,
|
||||
pending_updates=pending_updates,
|
||||
**transport_options,
|
||||
)
|
||||
return output
|
||||
|
||||
@@ -484,7 +626,7 @@ def unpack_diffusion_output_shm(output: object) -> object:
|
||||
|
||||
result = getattr(output, "result", None)
|
||||
if isinstance(result, DiffusionOutput):
|
||||
output.result = _unpack_diffusion_fields(result)
|
||||
setattr(output, "result", _unpack_diffusion_fields(result))
|
||||
|
||||
runner_outputs = getattr(output, "runner_outputs", None)
|
||||
if isinstance(runner_outputs, list):
|
||||
|
||||
@@ -2686,6 +2686,7 @@ class MiniMaxH3Pipeline(
|
||||
audio = audios[0] if len(audios) == 1 else torch.cat(audios, dim=0)
|
||||
return DiffusionOutput(
|
||||
output=(video, audio),
|
||||
video_output_index=0 if isinstance(video, torch.Tensor) else None,
|
||||
post_process_func=get_minimax_h3_post_process_func(self.od_config),
|
||||
stage_durations=(self.stage_durations if hasattr(self, "_stage_durations") else {}),
|
||||
)
|
||||
@@ -3049,6 +3050,7 @@ class MiniMaxH3Pipeline(
|
||||
self._release_stage_cache()
|
||||
return DiffusionOutput(
|
||||
output=(video, audio),
|
||||
video_output_index=0 if isinstance(video, torch.Tensor) else None,
|
||||
post_process_func=get_minimax_h3_post_process_func(self.od_config),
|
||||
stage_durations=(self.stage_durations if hasattr(self, "_stage_durations") else {}),
|
||||
)
|
||||
|
||||
@@ -1267,6 +1267,9 @@ class WorkerProc:
|
||||
"""
|
||||
device = torch.device(torch.accelerator.current_accelerator().type, self.gpu_id)
|
||||
d2h_stream = torch.Stream(device=device)
|
||||
transport_options = {}
|
||||
if self.od_config.video_output_transport.enable_registered_shm is True:
|
||||
transport_options["enable_registered_shm"] = True
|
||||
while True:
|
||||
item = self._async_output_queue.get()
|
||||
if item is None:
|
||||
@@ -1277,7 +1280,7 @@ class WorkerProc:
|
||||
# writing the output tensors before the side stream reads.
|
||||
if gpu_event is not None:
|
||||
d2h_stream.wait_event(gpu_event)
|
||||
pack_diffusion_output_shm(output, d2h_stream=d2h_stream)
|
||||
pack_diffusion_output_shm(output, d2h_stream=d2h_stream, **transport_options)
|
||||
d2h_stream.synchronize()
|
||||
|
||||
self._enqueue_result(
|
||||
|
||||
@@ -51,6 +51,7 @@ def split_diffusion_output_by_request(
|
||||
|
||||
return [
|
||||
DiffusionOutput(
|
||||
video_output_index=result.video_output_index,
|
||||
output=_slice_request_output(
|
||||
result.output,
|
||||
idx * num_outputs_per_prompt,
|
||||
|
||||
@@ -791,7 +791,10 @@ class OmniServeCommand(CLISubcommand):
|
||||
type=_json_object,
|
||||
default=None,
|
||||
help=(
|
||||
"JSON object configuring video output preparation, for example '{\"enable_device_postprocess\": true}'."
|
||||
"JSON object configuring video output transport: enable_device_postprocess enables GPU preparation; "
|
||||
"enable_registered_shm enables experimental CUDA-registered SHM for supported video outputs. "
|
||||
"enable_borrowed_frames enables borrowed RGB frames in MP4 response encoding. "
|
||||
"All default to false."
|
||||
),
|
||||
)
|
||||
omni_config_group.add_argument(
|
||||
|
||||
@@ -33,7 +33,7 @@ from vllm_omni.entrypoints.openai.stage_params import (
|
||||
build_stage_sampling_params_list,
|
||||
get_default_sampling_params_list,
|
||||
)
|
||||
from vllm_omni.entrypoints.openai.utils import is_video_generation_pipeline, parse_lora_request
|
||||
from vllm_omni.entrypoints.openai.utils import get_stage_type, is_video_generation_pipeline, parse_lora_request
|
||||
from vllm_omni.entrypoints.openai.video_api_utils import (
|
||||
_encode_video_bytes,
|
||||
_PlanarFrameConverter,
|
||||
@@ -187,6 +187,24 @@ class OmniOpenAIServingVideo:
|
||||
model_archs.extend(_stage_diffusion_model_class_name(stage) for stage in self.stage_configs or ())
|
||||
return od_config, tuple(get_diffusion_model_metadata(model_arch) for model_arch in model_archs)
|
||||
|
||||
def _video_encoding_options(self) -> dict[str, bool]:
|
||||
transport = _config_value(self._resolve_diffusion_od_config(), "video_output_transport")
|
||||
if transport is None:
|
||||
# Remote diffusion clients expose only model metadata to the API.
|
||||
# The resolved stage config still carries the effective transport options.
|
||||
stages = self.stage_configs or getattr(self._engine_client, "stage_configs", None) or ()
|
||||
for stage in stages:
|
||||
if get_stage_type(stage) != "diffusion":
|
||||
continue
|
||||
transport = _config_value(_config_value(stage, "diffusion_config"), "video_output_transport")
|
||||
if transport is None:
|
||||
transport = _config_value(_config_value(stage, "engine_args"), "video_output_transport")
|
||||
if transport is not None:
|
||||
break
|
||||
if _config_value(transport, "enable_borrowed_frames", False) is True:
|
||||
return {"enable_borrowed_frames": True}
|
||||
return {}
|
||||
|
||||
def _resolve_video_generation_defaults(
|
||||
self,
|
||||
request: VideoGenerationRequest,
|
||||
@@ -535,6 +553,8 @@ class OmniOpenAIServingVideo:
|
||||
if "video_codec_options" in request.extra_params:
|
||||
video_codec_options = request.extra_params["video_codec_options"]
|
||||
|
||||
encoding_options = self._video_encoding_options()
|
||||
|
||||
def encode_video_result(idx: int, video: Any) -> str:
|
||||
if isinstance(video, bytes):
|
||||
return base64.b64encode(video).decode("utf-8")
|
||||
@@ -545,6 +565,7 @@ class OmniOpenAIServingVideo:
|
||||
audio_sample_rate=artifacts.audio_sample_rate,
|
||||
video_codec_options=video_codec_options,
|
||||
frame_converter=self._video_frame_converter,
|
||||
**encoding_options,
|
||||
)
|
||||
|
||||
_t_encode_start = time.perf_counter()
|
||||
@@ -604,6 +625,9 @@ class OmniOpenAIServingVideo:
|
||||
logger.info("Action-only video request %s completed; skipping MP4 encoding.", reference_id)
|
||||
return b"", artifacts.stage_durations, artifacts.peak_memory_mb, action, video_metadata
|
||||
|
||||
encoding_options: dict[str, Any] = self._video_encoding_options()
|
||||
if audio is not None:
|
||||
encoding_options.update(audio=audio, audio_sample_rate=artifacts.audio_sample_rate)
|
||||
_t_encode_start = time.perf_counter()
|
||||
if isinstance(artifacts.videos[0], bytes):
|
||||
video_bytes = artifacts.videos[0]
|
||||
@@ -619,9 +643,9 @@ class OmniOpenAIServingVideo:
|
||||
video_bytes = _encode_video_bytes(
|
||||
artifacts.videos[0],
|
||||
fps=artifacts.output_fps,
|
||||
**({"audio": audio, "audio_sample_rate": artifacts.audio_sample_rate} if audio is not None else {}),
|
||||
video_codec_options=video_codec_options,
|
||||
frame_converter=self._video_frame_converter,
|
||||
**encoding_options,
|
||||
)
|
||||
_t_encode_ms = (time.perf_counter() - _t_encode_start) * 1000
|
||||
logger.info("Video response encoding (MP4 bytes): %.2f ms", _t_encode_ms)
|
||||
|
||||
@@ -749,13 +749,21 @@ def _iter_planar_video_frames(
|
||||
converter.shutdown()
|
||||
|
||||
|
||||
def _iter_borrowed_rgb_video_frames(frames: list[np.ndarray]) -> Generator[av.VideoFrame, None, None]:
|
||||
"""Wrap RGB arrays in AVFrames without copying their pixel storage."""
|
||||
import av
|
||||
|
||||
for frame in frames:
|
||||
yield av.VideoFrame.from_numpy_buffer(frame, format="rgb24")
|
||||
|
||||
|
||||
def _log_video_encoding_path(
|
||||
*,
|
||||
selected_path: str,
|
||||
frames: list[np.ndarray],
|
||||
frame_shape: tuple[int, ...],
|
||||
common_dtype: np.dtype,
|
||||
fps: int,
|
||||
fps: int | float,
|
||||
audio: AudioInput | None,
|
||||
audio_sample_rate: int | None,
|
||||
effective_frame_conversion_workers: int,
|
||||
@@ -790,7 +798,7 @@ def _encode_prepared_video_bytes_legacy(
|
||||
frames: list[np.ndarray],
|
||||
frame_shape: tuple[int, ...],
|
||||
common_dtype: np.dtype,
|
||||
fps: int,
|
||||
fps: int | float,
|
||||
audio: Any | None = None,
|
||||
audio_sample_rate: int | None = None,
|
||||
video_codec_options: dict[str, str] | None = None,
|
||||
@@ -835,14 +843,42 @@ def _encode_video_bytes(
|
||||
audio_sample_rate: int | None = None,
|
||||
video_codec_options: dict[str, str] | None = None,
|
||||
frame_converter: _PlanarFrameConverter | None = None,
|
||||
enable_borrowed_frames: bool = False,
|
||||
) -> bytes:
|
||||
"""Encode a video payload through the direct planar or legacy path."""
|
||||
"""Encode through the opt-in borrowed RGB, direct planar or legacy path."""
|
||||
from vllm_omni.diffusion.utils.media_utils import mux_av_video_audio_bytes
|
||||
|
||||
# Prepare once so validation is shared by both paths and malformed common
|
||||
# input is reported before any muxer is opened.
|
||||
frames, frame_shape, common_dtype = _prepare_video_frames(video)
|
||||
effective_audio_sample_rate = _resolve_audio_sample_rate(audio, audio_sample_rate) if audio is not None else None
|
||||
if (
|
||||
enable_borrowed_frames
|
||||
and common_dtype == np.dtype(np.uint8)
|
||||
and len(frame_shape) == 3
|
||||
and frame_shape[-1] == 3
|
||||
and all(frame.dtype == np.uint8 and frame.strides[1:] == (3, 1) for frame in frames)
|
||||
):
|
||||
_log_video_encoding_path(
|
||||
selected_path="borrowed_rgb",
|
||||
frames=frames,
|
||||
frame_shape=frame_shape,
|
||||
common_dtype=common_dtype,
|
||||
fps=fps,
|
||||
audio=audio,
|
||||
audio_sample_rate=effective_audio_sample_rate,
|
||||
effective_frame_conversion_workers=0,
|
||||
)
|
||||
audio_np = _coerce_audio_to_numpy(audio) if audio is not None else None
|
||||
return mux_av_video_audio_bytes(
|
||||
_iter_borrowed_rgb_video_frames(frames),
|
||||
width=frame_shape[1],
|
||||
height=frame_shape[0],
|
||||
audio_waveform=audio_np,
|
||||
fps=float(fps),
|
||||
audio_sample_rate=effective_audio_sample_rate,
|
||||
video_codec_options=video_codec_options,
|
||||
)
|
||||
fallback_reason = _direct_planar_fallback_reason(
|
||||
frames,
|
||||
frame_shape,
|
||||
@@ -957,6 +993,7 @@ def encode_video_base64(
|
||||
audio_sample_rate: int | None = None,
|
||||
video_codec_options: dict[str, str] | None = None,
|
||||
frame_converter: _PlanarFrameConverter | None = None,
|
||||
enable_borrowed_frames: bool = False,
|
||||
) -> str:
|
||||
"""Encode a video (frames/array/tensor) to base64 MP4."""
|
||||
video_bytes = _encode_video_bytes(
|
||||
@@ -966,5 +1003,6 @@ def encode_video_base64(
|
||||
audio_sample_rate=audio_sample_rate,
|
||||
video_codec_options=video_codec_options,
|
||||
frame_converter=frame_converter,
|
||||
**({"enable_borrowed_frames": True} if enable_borrowed_frames else {}),
|
||||
)
|
||||
return base64.b64encode(video_bytes).decode("utf-8")
|
||||
Reference in new issue
Block a user