diff --git a/docs/user_guide/diffusion/device_side_video_postprocess.md b/docs/user_guide/diffusion/device_side_video_postprocess.md index 910e64457..9f66702e4 100644 --- a/docs/user_guide/diffusion/device_side_video_postprocess.md +++ b/docs/user_guide/diffusion/device_side_video_postprocess.md @@ -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. diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 0dead804f..740961f81 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -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() diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 5dc39adb8..5e54637c2 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -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__}") diff --git a/vllm_omni/diffusion/ipc.py b/vllm_omni/diffusion/ipc.py index f0e88a2c9..88e069bbb 100644 --- a/vllm_omni/diffusion/ipc.py +++ b/vllm_omni/diffusion/ipc.py @@ -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): diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index a48fdc385..d08a5d255 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -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 {}), ) diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index f43e66906..fe344e4d0 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -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( diff --git a/vllm_omni/diffusion/worker/request_batch.py b/vllm_omni/diffusion/worker/request_batch.py index d5f1d5585..90d9abae5 100644 --- a/vllm_omni/diffusion/worker/request_batch.py +++ b/vllm_omni/diffusion/worker/request_batch.py @@ -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, diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index c26ba9375..ac1c3cda1 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -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( diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index c90a4d162..d4ecf20a4 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/vllm_omni/entrypoints/openai/serving_video.py @@ -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) diff --git a/vllm_omni/entrypoints/openai/video_api_utils.py b/vllm_omni/entrypoints/openai/video_api_utils.py index 67e8bcd36..43439913f 100644 --- a/vllm_omni/entrypoints/openai/video_api_utils.py +++ b/vllm_omni/entrypoints/openai/video_api_utils.py @@ -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")