diff --git a/docs/serving/videos_api.md b/docs/serving/videos_api.md index 09f023d27..c6e91df60 100644 --- a/docs/serving/videos_api.md +++ b/docs/serving/videos_api.md @@ -234,6 +234,34 @@ curl -X POST http://localhost:8091/v1/videos/sync \ -o output.mp4 ``` +## Output Encoding + +These `extra_params` control how the server turns decoded frames into MP4 bytes. + +| Field | Type | Default | Description | +| --- | --- | --- | --- | +| `preencode_mp4` | boolean | false | Encode the MP4 on the worker while the VAE is still decoding, instead of after the full video is materialized | +| `video_codec_options` | object | null | Encoder options passed through to the H.264 encoder, such as `{"preset": "ultrafast", "threads": "0"}` | + +With `preencode_mp4` enabled, each committed VAE chunk leaves the accelerator and +is encoded while later chunks are still decoding, so host transfer and CPU +encoding overlap the remaining decode instead of following it. The response is +unchanged: the same complete MP4, byte-for-byte equivalent frames. + +```bash +curl -X POST http://localhost:8091/v1/videos/sync \ + -F "prompt=A small robot walking through a neon city" \ + -F 'extra_params={"preencode_mp4": true, "video_codec_options": {"preset": "ultrafast"}}' \ + -o output.mp4 +``` + +`preencode_mp4` applies to the complete-MP4 response paths only. The +`/v1/realtime/video` WebSocket endpoint rejects it, because that path already +overlaps encoding through its own incremental fragmented-MP4 encoder. + +Support is per model: MiniMax-H3 implements it, and other models ignore the flag +and take the full-decode path. + ## Storage Set `VLLM_OMNI_SERVER_STORAGE__PATH` to control where asynchronous video outputs are diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index 5470cd9c4..7868c255b 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -34,6 +34,169 @@ def _turbo_spec(filename: str) -> TurboSpec: return spec +def test_decode_to_mp4_batches_consumer_transfers(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod + + class FakeEncoder: + instances = [] + + def __init__(self, **kwargs): + self.pushes = [] + self.kwargs = kwargs + self.__class__.instances.append(self) + + def push(self, frames): + self.pushes.append(np.array(frames, copy=True)) + + def finish(self): + return b"mp4" + + def abort(self): + raise AssertionError("unexpected abort") + + class FakeAudioVAE: + def decode_latent(self, latent): + return torch.zeros(1, 1, 2) + + class FakeVideoVAE: + def decode_with_chunks(self, latent, *, on_chunk): + for value in (0.0, 0.25, 0.5): + on_chunk(torch.full((1, 3, 3, 2, 2), value)) + + pipeline = object.__new__(mod.MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.audio_vae = FakeAudioVAE() + pipeline.video_vae = FakeVideoVAE() + pipeline.device = torch.device("cpu") + monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_uses_manual_component_offload", lambda self, component: False) + monkeypatch.setattr( + "vllm_omni.diffusion.utils.media_utils.ChunkedMP4Encoder", + FakeEncoder, + ) + + output = pipeline.decode_to_mp4( + torch.zeros(1), + torch.zeros(1), + height=2, + width=2, + batch_frames=4, + ) + + assert output == b"mp4" + encoder = FakeEncoder.instances[-1] + assert [frames.shape[0] for frames in encoder.pushes] == [6, 3] + assert np.array_equal(encoder.pushes[0][0], np.full((2, 2, 3), 0, dtype=np.uint8)) + assert np.array_equal(encoder.pushes[0][3], np.full((2, 2, 3), 64, dtype=np.uint8)) + assert encoder.pushes[1].shape == (3, 2, 2, 3) + + +def test_request_video_codec_options_reach_the_preencoded_mp4_encoder(monkeypatch): + """A client's encoder options must survive the worker-side pre-encode path.""" + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod + + class FakeEncoder: + instances = [] + + def __init__(self, **kwargs): + self.kwargs = kwargs + self.__class__.instances.append(self) + + def push(self, frames): + del frames + + def finish(self): + return b"mp4" + + def abort(self): + raise AssertionError("unexpected abort") + + class FakeAudioVAE: + def decode_latent(self, latent): + return torch.zeros(1, 1, 2) + + class FakeVideoVAE: + def decode_with_chunks(self, latent, *, on_chunk): + on_chunk(torch.zeros(1, 3, 1, 2, 2)) + + pipeline = object.__new__(mod.MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.audio_vae = FakeAudioVAE() + pipeline.video_vae = FakeVideoVAE() + pipeline.device = torch.device("cpu") + monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_uses_manual_component_offload", lambda self, component: False) + monkeypatch.setattr("vllm_omni.diffusion.utils.media_utils.ChunkedMP4Encoder", FakeEncoder) + + pipeline.decode_to_mp4( + torch.zeros(1), + torch.zeros(1), + height=2, + width=2, + video_codec_options={"preset": "ultrafast"}, + ) + + assert FakeEncoder.instances[-1].kwargs["video_codec_options"] == {"preset": "ultrafast"} + + +@pytest.mark.parametrize( + ("codec_extra", "expected"), + [ + ({}, {"preset": "ultrafast", "threads": "0"}), + ({"video_codec_options": {"preset": "slow", "threads": 2}}, {"preset": "slow", "threads": "2"}), + ({"video_codec_options": {}}, {}), + ({"video_codec_options": None}, None), + ], +) +@pytest.mark.parametrize("batch_extra, batch_frames", [({}, 17), ({"preencode_batch_frames": 5}, 5)]) +def test_preencode_request_preserves_serving_codec_defaults(codec_extra, expected, batch_extra, batch_frames): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "fl2va" + pipeline.supported_tasks = frozenset({"t2va"}) + pipeline.default_video_shift = 12.0 + pipeline.default_audio_shift = 3.0 + pipeline.device = torch.device("cpu") + pipeline.od_config = SimpleNamespace() + pipeline.text_encoder = object() + pipeline.encode_prompt = Mock(return_value=(torch.ones(1, 2), torch.ones(1, dtype=torch.long))) + pipeline._quality_policy = Mock() + pipeline._cache_dit_runtime = Mock() + pipeline.diffuse = Mock(return_value=(torch.zeros(1), torch.zeros(1))) + pipeline.decode_to_mp4 = Mock(return_value=b"mp4") + sampling = OmniDiffusionSamplingParams( + width=1344, + height=768, + fps=24, + num_frames=124, + num_inference_steps=50, + extra_args={"task": "t2va", "aspect_ratio": "16:9", "preencode_mp4": True, **codec_extra, **batch_extra}, + ) + batch = DiffusionRequestBatch( + [OmniDiffusionRequest(prompt="test", sampling_params=sampling, request_id="codec-defaults")] + ) + + pipeline.forward(batch) + + assert pipeline.decode_to_mp4.call_args.kwargs["video_codec_options"] == expected + assert pipeline.decode_to_mp4.call_args.kwargs["batch_frames"] == batch_frames + + +def test_video_codec_options_are_normalized_for_the_encoder(): + from vllm_omni.diffusion.utils.media_utils import normalize_video_codec_options + + assert normalize_video_codec_options({"preset": "ultrafast", "threads": 0}) == { + "preset": "ultrafast", + "threads": "0", + } + assert normalize_video_codec_options(None) is None + with pytest.raises(ValueError, match="video_codec_options"): + normalize_video_codec_options("preset=ultrafast") + + def test_h3_prepares_resolved_cache_state_immediately_before_denoise(): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline from vllm_omni.diffusion.request import OmniDiffusionRequest diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py index 63c3a42f1..1b123929c 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py @@ -284,7 +284,8 @@ def test_mixed_step_batch_leaves_gated_attention_dense(): assert recorder.denoise_timestep is None -def test_prepare_encode_seeds_runner_visible_state(monkeypatch): +@pytest.mark.parametrize("batch_frames", [1, 33]) +def test_prepare_encode_seeds_runner_visible_state(monkeypatch, batch_frames): from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod branch, video_rows, audio_rows = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=8) @@ -293,6 +294,8 @@ def test_prepare_encode_seeds_runner_visible_state(monkeypatch): context = { "height": 96, "width": 64, + "preencode_mp4": True, + "preencode_batch_frames": batch_frames, "latent_t": 2, "latent_h": 4, "latent_w": 6, @@ -335,6 +338,18 @@ def test_prepare_encode_seeds_runner_visible_state(monkeypatch): assert state.extra[mod._STEP_BRANCH] is branch assert state.extra[mod._STEP_SHAPE]["height"] == 96 + pipeline.od_config = SimpleNamespace() + monkeypatch.setattr(pipeline, "_unpack_denoised_rows", lambda *args, **kwargs: (torch.zeros(1), torch.zeros(1))) + calls = [] + + def decode_to_mp4(*args, **kwargs): + calls.append(kwargs) + return b"mp4" + + monkeypatch.setattr(pipeline, "decode_to_mp4", decode_to_mp4) + assert pipeline.post_decode(state).output == (b"mp4", None) + assert calls[0]["batch_frames"] == batch_frames + def test_prepare_encode_rejects_request_mode_only_features(): """Multi-output and DLO have no representation in the step contract.""" diff --git a/tests/entrypoints/openai_api/test_serving_video_output_stream.py b/tests/entrypoints/openai_api/test_serving_video_output_stream.py index e57e3e279..291eada75 100644 --- a/tests/entrypoints/openai_api/test_serving_video_output_stream.py +++ b/tests/entrypoints/openai_api/test_serving_video_output_stream.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for `/v1/realtime/video` WebSocket video output streaming.""" from __future__ import annotations @@ -7,12 +7,13 @@ from __future__ import annotations import asyncio import json from collections.abc import AsyncGenerator, Callable +from http import HTTPStatus from types import SimpleNamespace from unittest.mock import MagicMock import numpy as np import pytest -from fastapi import FastAPI, WebSocket +from fastapi import FastAPI, HTTPException, WebSocket from pytest_mock import MockerFixture from starlette.testclient import TestClient @@ -957,3 +958,30 @@ class TestStreamingVideoOutputPromptUpdate: "transition_chunks": 2, }, ) + + +class TestStreamingVideoOutputExtraParams: + """extra_params handling when building streaming generation inputs.""" + + @staticmethod + def _handler() -> OmniStreamingVideoOutputHandler: + handler = object.__new__(OmniStreamingVideoOutputHandler) + handler._engine_client = SimpleNamespace(default_sampling_params_list=[]) + return handler + + async def test_preencode_mp4_is_rejected_for_streaming_sessions(self): + """Worker-side pre-encoding emits one progressive MP4 the fMP4 encoder cannot consume.""" + request = VideoGenerationRequest(prompt="p", extra_params={"preencode_mp4": True}) + + with pytest.raises(HTTPException) as excinfo: + await self._handler()._build_prompt_and_sampling_params(request) + + assert excinfo.value.status_code == HTTPStatus.BAD_REQUEST.value + assert "preencode_mp4" in excinfo.value.detail + + async def test_other_extra_params_still_reach_the_engine(self): + request = VideoGenerationRequest(prompt="p", extra_params={"flow_shift": 3.0}) + + _, gen_params, _ = await self._handler()._build_prompt_and_sampling_params(request) + + assert gen_params.extra_args["flow_shift"] == 3.0 diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index b65b00e59..01332beec 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -134,6 +134,60 @@ def test_raw_and_base64_encoders_receive_persistent_converter(mocker: MockerFixt 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() + handler = OmniOpenAIServingVideo.for_diffusion(engine, model_name="test-model") + request = VideoGenerationRequest( + prompt="test", extra_params={"preencode_mp4": True, "preencode_batch_frames": batch_frames} + ) + try: + with pytest.raises(HTTPException, match="preencode_batch_frames") as exc: + asyncio.run(handler.generate_video_bytes(request, "invalid-batch")) + assert exc.value.status_code == 400 + assert engine.captured_prompt is None + finally: + handler.shutdown() + + +def test_preencoded_video_bytes_preserve_metadata(mocker: MockerFixture): + from vllm_omni.entrypoints.openai.serving_video import VideoGenerationArtifacts + + handler = OmniOpenAIServingVideo.for_diffusion(FakeAsyncOmni(), model_name="test-model") + # Resolved frame count differs from anything the request asked for, so the + # metadata has to come from the encoded stream rather than request defaults. + preencoded = _make_test_video_bytes((32, 24), num_frames=7) + artifacts = VideoGenerationArtifacts( + videos=[preencoded], + audios=[None], + actions=[None], + audio_sample_rate=24000, + output_fps=24.0, + stage_durations={"decode": 0.5}, + peak_memory_mb=123.0, + metrics={"generation_time": 1.25}, + ) + mocker.patch.object(handler, "_run_and_extract", return_value=artifacts) + encoder = mocker.patch("vllm_omni.entrypoints.openai.serving_video._encode_video_bytes") + try: + result = asyncio.run(handler.generate_video_bytes(VideoGenerationRequest(prompt="test"), "preencoded")) + assert result == ( + preencoded, + {"decode": 0.5}, + 123.0, + None, + { + "fps": 24.0, + "num_frames": 7, + "duration_s": 7 / 24.0, + "metrics": {"generation_time": 1.25}, + }, + ) + encoder.assert_not_called() + finally: + handler.shutdown() + + def test_resolve_diffusion_od_config_falls_back_to_attribute(): od_config = SimpleNamespace(model_class_name="WanPipeline") handler = OmniOpenAIServingVideo.for_diffusion( 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 fb9f02b28..f80513d4e 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -64,6 +64,7 @@ from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import ( DiffusionPipelineProfilerMixin, ) from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule +from vllm_omni.diffusion.utils.media_utils import normalize_preencode_batch_frames, normalize_video_codec_options from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.errors import OmniClientError, client_error_from_metadata from vllm_omni.model_executor.model_loader.weight_utils import ( @@ -367,6 +368,19 @@ def _minimax_h3_post_process(output, output_type: str = "np"): if not isinstance(output, tuple) or len(output) != 2: return output video, audio = output + if isinstance(video, (bytes, bytearray, memoryview)): + video = [video] + if isinstance(video, list) and all(isinstance(item, (bytes, bytearray, memoryview)) for item in video): + encoded_videos = [bytes(item) for item in video] + else: + encoded_videos = None + if encoded_videos is not None: + return { + "video": encoded_videos, + "audio": [None] * len(encoded_videos), + "audio_sample_rate": MINIMAX_H3_AUDIO_SAMPLE_RATE, + "fps": MINIMAX_H3_FPS, + } if video.dtype != torch.uint8 or video.ndim != 5 or video.shape[-1] not in (3, 4): # Float or channel-first frames would reach the muxer as a black or # banded video rather than as an error. @@ -2147,6 +2161,80 @@ class MiniMaxH3Pipeline( audio_t=audio_t, ) + def decode_to_mp4( + self, + video_latent: torch.Tensor, + audio_latent: torch.Tensor, + *, + height: int, + width: int, + max_pending: int = 2, + batch_frames: int = 17, + video_codec_options: dict[str, str] | None = None, + ) -> bytes: + """Decode and encode one output on the worker without full-video materialization. + + Audio is decoded first so the incremental mux session can attach its audio + stream before temporal video chunks arrive. The callback receives committed + float32 ``BCTHW`` frames, performs the requested-size crop and uint8 + conversion in the worker, then applies bounded backpressure to the encoder. + """ + from vllm_omni.diffusion.utils.media_utils import ChunkedMP4Encoder + + if batch_frames <= 0: + raise ValueError("batch_frames must be positive") + + with self._component_on_device(self.audio_vae): + audio = self.audio_vae.decode_latent(audio_latent) + audio_np = audio.detach().float().cpu().numpy() + if audio_np.ndim == 3 and audio_np.shape[0] == 1: + audio_np = audio_np[0] + encoder = ChunkedMP4Encoder( + width=width, + height=height, + fps=MINIMAX_H3_FPS, + audio_waveform=audio_np, + audio_sample_rate=MINIMAX_H3_AUDIO_SAMPLE_RATE, + max_pending=max_pending, + video_codec_options=video_codec_options, + ) + + pending_chunks: list[torch.Tensor] = [] + pending_frames = 0 + + def flush_pending() -> None: + nonlocal pending_frames + if not pending_chunks: + return + batched = torch.cat(pending_chunks, dim=1) + encoder.push(batched[0].cpu().numpy()) + pending_chunks.clear() + pending_frames = 0 + + def on_chunk(frames: torch.Tensor) -> None: + nonlocal pending_frames + prepared = _prepare_minimax_h3_video_output(frames[..., :height, :width]) + if prepared.shape[0] != 1: + raise ValueError("MiniMax H3 chunked MP4 encoding currently expects one output per decoder") + pending_chunks.append(prepared) + pending_frames += int(prepared.shape[1]) + if pending_frames >= batch_frames: + flush_pending() + + try: + with self._component_on_device(self.video_vae): + with current_omni_platform.create_autocast_context( + device_type=self.device.type, + dtype=torch.float16, + enabled=True, + ): + self.video_vae.decode_with_chunks(video_latent, on_chunk=on_chunk) + flush_pending() + return encoder.finish() + except BaseException: + encoder.abort() + raise + def decode( self, video_latent: torch.Tensor, @@ -2230,6 +2318,11 @@ class MiniMaxH3Pipeline( quality = sampling.quality logger.debug("MiniMax H3 request quality=%s", quality) extra = sampling.extra_args or {} + preencode_batch_frames = ( + normalize_preencode_batch_frames(extra.get("preencode_batch_frames", 17)) + if extra.get("preencode_mp4", False) + else 17 + ) turbo_spec = self._active_turbo_spec(sampling) has_native_lora = self._has_active_native_lora(sampling) task = self._resolve_task( @@ -2500,6 +2593,11 @@ class MiniMaxH3Pipeline( "audio_shift": audio_shift, "base_schedule": base_schedule, "num_outputs": num_outputs, + "preencode_mp4": bool(extra.get("preencode_mp4", False)), + "preencode_batch_frames": preencode_batch_frames, + "video_codec_options": normalize_video_codec_options( + extra.get("video_codec_options", {"preset": "ultrafast", "threads": "0"}) + ), } @staticmethod @@ -2526,16 +2624,33 @@ class MiniMaxH3Pipeline( audios = [] for output_seed in _minimax_h3_output_seeds(context["seed"], num_outputs): video_latent, audio_latent = self.diffuse(**{**denoise_kwargs, "seed": output_seed}) - video, audio = self.decode( - video_latent, - audio_latent, - height=context["height"], - width=context["width"], - ) - videos.append(_prepare_minimax_h3_video_output(video)) - audios.append(audio) - video = videos[0] if len(videos) == 1 else torch.cat(videos, dim=0) - audio = audios[0] if len(audios) == 1 else torch.cat(audios, dim=0) + if context["preencode_mp4"]: + videos.append( + self.decode_to_mp4( + video_latent, + audio_latent, + height=context["height"], + width=context["width"], + video_codec_options=context["video_codec_options"], + batch_frames=context["preencode_batch_frames"], + ) + ) + audios.append(None) + else: + video, audio = self.decode( + video_latent, + audio_latent, + height=context["height"], + width=context["width"], + ) + videos.append(_prepare_minimax_h3_video_output(video)) + audios.append(audio) + if videos and isinstance(videos[0], bytes): + video = videos[0] if len(videos) == 1 else videos + audio = None + else: + video = videos[0] if len(videos) == 1 else torch.cat(videos, dim=0) + audio = audios[0] if len(audios) == 1 else torch.cat(audios, dim=0) return DiffusionOutput( output=(video, audio), post_process_func=get_minimax_h3_post_process_func(self.od_config), @@ -2660,6 +2775,9 @@ class MiniMaxH3Pipeline( "latent_h": context["latent_h"], "latent_w": context["latent_w"], "audio_t": context["audio_t"], + "preencode_mp4": context.get("preencode_mp4", False), + "preencode_batch_frames": context.get("preencode_batch_frames", 17), + "video_codec_options": context.get("video_codec_options"), }, } ) @@ -2831,13 +2949,24 @@ class MiniMaxH3Pipeline( latent_w=shape["latent_w"], audio_t=shape["audio_t"], ) - video, audio = self.decode( - video_latent, - audio_latent, - height=shape["height"], - width=shape["width"], - ) - video = _prepare_minimax_h3_video_output(video) + if shape.get("preencode_mp4", False): + video = self.decode_to_mp4( + video_latent, + audio_latent, + height=shape["height"], + width=shape["width"], + video_codec_options=shape.get("video_codec_options"), + batch_frames=shape.get("preencode_batch_frames", 17), + ) + audio = None + else: + video, audio = self.decode( + video_latent, + audio_latent, + height=shape["height"], + width=shape["width"], + ) + video = _prepare_minimax_h3_video_output(video) return DiffusionOutput( output=(video, audio), post_process_func=get_minimax_h3_post_process_func(self.od_config), diff --git a/vllm_omni/diffusion/utils/media_utils.py b/vllm_omni/diffusion/utils/media_utils.py index 615df090f..99d20a0af 100644 --- a/vllm_omni/diffusion/utils/media_utils.py +++ b/vllm_omni/diffusion/utils/media_utils.py @@ -7,7 +7,7 @@ from __future__ import annotations import io import queue import threading -from collections.abc import Iterable +from collections.abc import Iterable, Mapping from fractions import Fraction from typing import Any, cast @@ -17,6 +17,27 @@ import numpy as np _CHUNKED_MP4_DONE = object() +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: @@ -233,6 +254,18 @@ class FragmentedMP4Muxer: 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, *, diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index b3afe8b40..5ff055fa2 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/vllm_omni/entrypoints/openai/serving_video.py @@ -13,12 +13,14 @@ from http import HTTPStatus from types import SimpleNamespace from typing import TYPE_CHECKING, Any, cast +import pybase64 as base64 from fastapi import HTTPException from PIL import Image from vllm.engine.protocol import EngineClient from vllm.logger import init_logger from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata +from vllm_omni.diffusion.utils.media_utils import count_mp4_frames, normalize_preencode_batch_frames from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.entrypoints.openai.protocol.videos import ( VideoAction, @@ -99,7 +101,12 @@ def _video_metadata_from_artifacts(artifacts: VideoGenerationArtifacts) -> dict[ if not artifacts.videos: return metadata - num_frames = count_video_frames(artifacts.videos[0]) + video = artifacts.videos[0] + # Pre-encoded outputs arrive as MP4 bytes, which carry no tensor shape. + if isinstance(video, (bytes, bytearray, memoryview)): + num_frames = count_mp4_frames(bytes(video)) + else: + num_frames = count_video_frames(video) if num_frames is not None and num_frames > 0: metadata["num_frames"] = num_frames if artifacts.output_fps > 0: @@ -402,6 +409,11 @@ class OmniOpenAIServingVideo: status_code=HTTPStatus.BAD_REQUEST.value, detail="extra_params must be a JSON object/dict.", ) + if request.extra_params.get("preencode_mp4") and "preencode_batch_frames" in request.extra_params: + try: + normalize_preencode_batch_frames(request.extra_params["preencode_batch_frames"]) + except ValueError as exc: + raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(exc)) from exc # Merge extra_params into extra_args gen_params.extra_args.update(request.extra_params) @@ -479,26 +491,22 @@ class OmniOpenAIServingVideo: if "video_codec_options" in request.extra_params: video_codec_options = request.extra_params["video_codec_options"] + def encode_video_result(idx: int, video: Any) -> str: + if isinstance(video, bytes): + return base64.b64encode(video).decode("utf-8") + return encode_video_base64( + video, + fps=artifacts.output_fps, + audio=artifacts.audios[idx], + audio_sample_rate=artifacts.audio_sample_rate, + video_codec_options=video_codec_options, + frame_converter=self._video_frame_converter, + ) + _t_encode_start = time.perf_counter() video_data = [ VideoData( - b64_json=( - encode_video_base64( - video, - fps=artifacts.output_fps, - video_codec_options=video_codec_options, - frame_converter=self._video_frame_converter, - ) - if artifacts.audios[idx] is None - else encode_video_base64( - video, - fps=artifacts.output_fps, - audio=artifacts.audios[idx], - audio_sample_rate=artifacts.audio_sample_rate, - video_codec_options=video_codec_options, - frame_converter=self._video_frame_converter, - ) - ), + b64_json=encode_video_result(idx, video), action=artifacts.actions[idx], ) for idx, video in enumerate(artifacts.videos) @@ -549,6 +557,17 @@ class OmniOpenAIServingVideo: return b"", artifacts.stage_durations, artifacts.peak_memory_mb, action, video_metadata _t_encode_start = time.perf_counter() + if isinstance(artifacts.videos[0], bytes): + video_bytes = artifacts.videos[0] + _t_encode_ms = (time.perf_counter() - _t_encode_start) * 1000 + logger.info("Video response received pre-encoded MP4 bytes: %.2f ms", _t_encode_ms) + return ( + video_bytes, + artifacts.stage_durations, + artifacts.peak_memory_mb, + artifacts.actions[0], + video_metadata, + ) video_bytes = _encode_video_bytes( artifacts.videos[0], fps=artifacts.output_fps, diff --git a/vllm_omni/entrypoints/openai/serving_video_output_stream.py b/vllm_omni/entrypoints/openai/serving_video_output_stream.py index a9abad146..f8f4722ad 100644 --- a/vllm_omni/entrypoints/openai/serving_video_output_stream.py +++ b/vllm_omni/entrypoints/openai/serving_video_output_stream.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """WebSocket handler for streaming generated video chunks. Protocol: @@ -556,6 +556,14 @@ class OmniStreamingVideoOutputHandler: status_code=HTTPStatus.BAD_REQUEST.value, detail="extra_params must be a JSON object/dict.", ) + # Worker-side pre-encoding produces one progressive MP4, which this + # path's incremental fMP4 encoder cannot consume; streaming already + # overlaps encoding through that encoder instead. + if request.extra_params.get("preencode_mp4"): + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="preencode_mp4 is not supported for streaming video sessions.", + ) gen_params.extra_args.update(request.extra_params) self._apply_lora(request.lora, gen_params)