diff --git a/docs/.nav.yml b/docs/.nav.yml index af80816e4..ddf40376c 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -62,6 +62,7 @@ nav: - FP8: user_guide/diffusion/quantization/fp8.md - Int8: user_guide/diffusion/quantization/int8.md - GGUF: user_guide/diffusion/quantization/gguf.md + - Frame Interpolation: user_guide/diffusion/frame_interpolation.md - Step Execution: user_guide/diffusion/step_execution.md - Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md diff --git a/docs/user_guide/diffusion/frame_interpolation.md b/docs/user_guide/diffusion/frame_interpolation.md new file mode 100644 index 000000000..349af50c5 --- /dev/null +++ b/docs/user_guide/diffusion/frame_interpolation.md @@ -0,0 +1,92 @@ +# Frame Interpolation + +## Overview + +vLLM-Omni supports post-generation frame interpolation for supported video +diffusion pipelines. This feature inserts synthesized intermediate frames +between adjacent generated frames to improve temporal smoothness without +rerunning the diffusion denoising loop. + +Frame interpolation runs in the diffusion worker post-processing path instead +of the API server encoding path. This allows the interpolation step to reuse +the worker's current accelerator device and keeps the FastAPI event loop free +from heavy synchronous PyTorch work. + +For an input video with `N` generated frames and interpolation exponent `exp`, +the output frame count is: + +```text +(N - 1) * 2**exp + 1 +``` + +The output FPS is multiplied by `2**exp` so the clip duration remains close to +the original generated video. + +## Supported Pipelines + +Frame interpolation is currently supported for: + +- `WanPipeline` (Wan2.2 text-to-video) +- `WanImageToVideoPipeline` +- `Wan22TI2VPipeline` + +## Request Parameters + +The video APIs `/v1/videos` and `/v1/videos/sync` accept: + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `enable_frame_interpolation` | bool | `false` | Enable post-generation frame interpolation | +| `frame_interpolation_exp` | int | `1` | Interpolation exponent. `1=2x`, `2=4x`, etc. | +| `frame_interpolation_scale` | float | `1.0` | RIFE inference scale | +| `frame_interpolation_model_path` | str | `None` | Local directory or Hugging Face repo ID containing `flownet.pkl` | + +## Execution Flow + +For supported Wan2.2 pipelines, the execution order is: + +1. Diffusion worker finishes denoising and decodes the raw video tensor. +2. Worker-side model-specific post-processing runs. +3. If frame interpolation is enabled, RIFE interpolates the decoded video + tensor on the worker side and records a FPS multiplier in `custom_output`. +4. The API server receives the already-interpolated video and only performs + MP4 export. + +This design keeps interpolation close to the generated tensor and avoids +introducing another heavyweight GPU context in the API server process. + +## Example + +Start the server: + +```bash +vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 +``` + +Run a sync request with interpolation enabled: + +```bash +curl -X POST http://localhost:8091/v1/videos/sync \ + -F "prompt=A dog running through a park" \ + -F "num_frames=81" \ + -F "width=832" \ + -F "height=480" \ + -F "fps=16" \ + -F "num_inference_steps=40" \ + -F "guidance_scale=1.0" \ + -F "guidance_scale_2=1.0" \ + -F "enable_frame_interpolation=true" \ + -F "frame_interpolation_exp=1" \ + -F "frame_interpolation_scale=1.0" \ + -F "seed=42" \ + -o sync_t2v_interpolated.mp4 +``` + +## Notes + +- This is a post-processing feature. It does not modify the diffusion denoising + schedule. +- Higher interpolation exponents increase post-processing time and memory usage. +- If the interpolation model weights are not available locally, + `frame_interpolation_model_path` may point to a Hugging Face repo containing + `flownet.pkl`. diff --git a/docs/user_guide/examples/online_serving/text_to_video.md b/docs/user_guide/examples/online_serving/text_to_video.md index ea7a94ed9..f49d3e318 100644 --- a/docs/user_guide/examples/online_serving/text_to_video.md +++ b/docs/user_guide/examples/online_serving/text_to_video.md @@ -123,6 +123,9 @@ curl -X POST http://localhost:8091/v1/videos \ -F "guidance_scale_2=4.0" \ -F "boundary_ratio=0.875" \ -F "flow_shift=5.0" \ + -F "enable_frame_interpolation=true" \ + -F "frame_interpolation_exp=1" \ + -F "frame_interpolation_scale=1.0" \ -F "seed=42" ``` @@ -145,6 +148,35 @@ curl -X POST http://localhost:8091/v1/videos \ | `flow_shift` | float | None | Scheduler flow shift (Wan2.2) | | `seed` | int | None | Random seed (reproducible) | | `lora` | object | None | LoRA configuration | +| `enable_frame_interpolation` | bool | false | Enable RIFE frame interpolation before MP4 encoding | +| `frame_interpolation_exp` | int | 1 | Interpolation exponent; 1=2x temporal resolution, 2=4x | +| `frame_interpolation_scale` | float | 1.0 | RIFE inference scale; use 0.5 for high-resolution inputs | +| `frame_interpolation_model_path` | str | None | Local directory or Hugging Face repo ID with `flownet.pkl`; defaults to `elfgum/RIFE-4.22.lite` | + +## Frame Interpolation + +Frame interpolation is an optional post-processing step for `/v1/videos` and +`/v1/videos/sync`. It synthesizes intermediate frames between generated frames +without rerunning the diffusion model. If the generated video has `N` frames, +the interpolated output frame count is `(N - 1) * 2**exp + 1`. The encoder FPS +is multiplied by `2**exp` so the output duration remains close to the original. + +Frame interpolation runs in the diffusion worker post-processing path instead of +the API server encoding path, so it can reuse the worker's current accelerator +device without blocking the FastAPI event loop. + +Example: generate 5 frames and interpolate to 9 frames: + +```bash +curl -X POST http://localhost:8091/v1/videos/sync \ + -F "prompt=A dog running through a park" \ + -F "num_frames=5" \ + -F "fps=8" \ + -F "enable_frame_interpolation=true" \ + -F "frame_interpolation_exp=1" \ + -F "frame_interpolation_scale=1.0" \ + -o sync_t2v_interpolated.mp4 +``` ## Create Response Format diff --git a/tests/diffusion/test_diffusion_step_pipeline.py b/tests/diffusion/test_diffusion_step_pipeline.py index ad08487fe..9b7dd9cc4 100644 --- a/tests/diffusion/test_diffusion_step_pipeline.py +++ b/tests/diffusion/test_diffusion_step_pipeline.py @@ -443,14 +443,42 @@ class TestWorker: class TestIPC: def test_pack_unpack_runner_output_shm(self): tensor = torch.zeros(300_000, dtype=torch.float32) - output = RunnerOutput(req_id="req-1", finished=True, result=DiffusionOutput(output=tensor)) + log_probs = torch.ones(300_000, dtype=torch.float32) + output = RunnerOutput( + req_id="req-1", + finished=True, + result=DiffusionOutput(output=tensor, trajectory_log_probs=log_probs), + ) packed = pack_diffusion_output_shm(output) assert isinstance(packed.result.output, dict) assert packed.result.output["__tensor_shm__"] is True + assert isinstance(packed.result.trajectory_log_probs, dict) + assert packed.result.trajectory_log_probs["__tensor_shm__"] is True unpacked = unpack_diffusion_output_shm(packed) assert isinstance(unpacked.result.output, torch.Tensor) + assert isinstance(unpacked.result.trajectory_log_probs, torch.Tensor) + torch.testing.assert_close(unpacked.result.output, tensor) + torch.testing.assert_close(unpacked.result.trajectory_log_probs, log_probs) + + def test_pack_unpack_runner_output_shm_bfloat16(self): + tensor = torch.arange(600_000, dtype=torch.bfloat16) + output = RunnerOutput( + req_id="req-1", + finished=True, + result=DiffusionOutput(output=tensor), + ) + + packed = pack_diffusion_output_shm(output) + assert isinstance(packed.result.output, dict) + assert packed.result.output["__tensor_shm__"] is True + assert packed.result.output["torch_dtype"] == str(torch.bfloat16) + assert packed.result.output["numpy_dtype"] == "uint16" + + unpacked = unpack_diffusion_output_shm(packed) + assert isinstance(unpacked.result.output, torch.Tensor) + assert unpacked.result.output.dtype == torch.bfloat16 torch.testing.assert_close(unpacked.result.output, tensor) diff --git a/tests/diffusion/test_inline_stage_diffusion_client.py b/tests/diffusion/test_inline_stage_diffusion_client.py new file mode 100644 index 000000000..5ea55499f --- /dev/null +++ b/tests/diffusion/test_inline_stage_diffusion_client.py @@ -0,0 +1,122 @@ +from __future__ import annotations + +import asyncio +from unittest.mock import MagicMock, patch + +import pytest + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient +from vllm_omni.engine.stage_init_utils import StageMetadata +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.outputs import OmniRequestOutput + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture +def mock_engine(): + with patch("vllm_omni.diffusion.inline_stage_diffusion_client.DiffusionEngine") as mock: + engine_instance = MagicMock() + mock.make_engine.return_value = engine_instance + yield engine_instance + + +@pytest.fixture +def client(mock_engine): + metadata = StageMetadata( + stage_id=0, + stage_type="diffusion", + engine_output_type="image", + is_comprehension=False, + requires_multimodal_data=False, + engine_input_source="prompt", + final_output=True, + final_output_type="image", + default_sampling_params={}, + custom_process_input_func=None, + model_stage=None, + runtime_cfg=None, + ) + with patch.object(InlineStageDiffusionClient, "_enrich_config"): + od_config = MagicMock(spec=OmniDiffusionConfig) + c = InlineStageDiffusionClient(model="test_model", od_config=od_config, metadata=metadata, batch_size=1) + yield c + c.shutdown() + + +@pytest.mark.asyncio +async def test_inline_dispatch_request_success(client, mock_engine): + mock_result = OmniRequestOutput.from_diffusion(request_id="req-1", images=[MagicMock()]) + mock_engine.step.return_value = [mock_result] + + sampling_params = OmniDiffusionSamplingParams() + await client.add_request_async("req-1", "A test prompt", sampling_params) + + # Wait for the task to be processed + for _ in range(10): + output = client.get_diffusion_output_nowait() + if output is not None: + break + await asyncio.sleep(0.01) + + assert output is not None + assert output.request_id == "req-1" + mock_engine.step.assert_called_once() + request = mock_engine.step.call_args.args[0] + assert request.prompts == ["A test prompt"] + assert request.request_ids == ["req-1"] + + +@pytest.mark.asyncio +async def test_inline_dispatch_batch_success(client, mock_engine): + mock_engine.step.return_value = [ + OmniRequestOutput.from_diffusion(request_id="req-batch-0", images=[MagicMock()]), + OmniRequestOutput.from_diffusion(request_id="req-batch-1", images=[MagicMock()]), + ] + + sampling_params = OmniDiffusionSamplingParams() + prompts = ["prompt-0", "prompt-1"] + await client.add_batch_request_async("req-batch", prompts, sampling_params) + + for _ in range(10): + output = client.get_diffusion_output_nowait() + if output is not None: + break + await asyncio.sleep(0.01) + + assert output is not None + assert output.request_id == "req-batch" + assert len(output.images) == 2 + request = mock_engine.step.call_args.args[0] + assert request.prompts == prompts + assert request.request_ids == ["req-batch-0", "req-batch-1"] + + +@pytest.mark.asyncio +async def test_inline_dispatch_request_error(client, mock_engine): + mock_engine.step.side_effect = RuntimeError("Engine failure") + + sampling_params = OmniDiffusionSamplingParams() + await client.add_request_async("req-err", "A test prompt", sampling_params) + + for _ in range(10): + output = client.get_diffusion_output_nowait() + if output is not None: + break + await asyncio.sleep(0.01) + + assert output is not None + assert output["request_id"] == "req-err" + assert output["error"] == "Engine failure" + assert output["error_type"] == "RuntimeError" + + +def test_inline_shutdown(client, mock_engine): + assert not client._shutting_down + + # Shutting down should cleanly cancel anything queued and close engine + client.shutdown() + + assert client._shutting_down + mock_engine.close.assert_called_once() diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 28b44e9bd..f95aa53af 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -70,6 +70,66 @@ def test_initialize_stages_restores_device_visibility_after_diffusion_init(monke else: os.environ[env_var] = old_env +def test_initialize_stages_uses_inline_diffusion_client_for_single_stage(monkeypatch): + """Single-stage diffusion init should request the inline client path.""" + import vllm_omni.engine.async_omni_engine as engine_mod + from vllm_omni.platforms import current_omni_platform + + engine = object.__new__(AsyncOmniEngine) + engine.model = "dummy-model" + engine.config_path = "dummy-config" + engine.num_stages = 1 + engine.async_chunk = False + engine.diffusion_batch_size = 1 + engine.stage_configs = [types.SimpleNamespace(stage_id=0, stage_type="diffusion")] + + env_var = current_omni_platform.device_control_env_var + old_env = os.environ.get(env_var) + os.environ[env_var] = "0" + + metadata = types.SimpleNamespace( + stage_id=0, + stage_type="diffusion", + runtime_cfg={"devices": "0"}, + prompt_expand_func=None, + ) + captured_use_inline = None + diffusion_client = types.SimpleNamespace(is_comprehension=False) + + monkeypatch.setattr(engine_mod, "prepare_engine_environment", lambda: None) + monkeypatch.setattr(engine_mod, "load_omni_transfer_config_for_model", lambda *_: None) + monkeypatch.setattr(engine_mod, "extract_stage_metadata", lambda _cfg: metadata) + monkeypatch.setattr(engine_mod, "get_stage_connector_spec", lambda **_: {}) + monkeypatch.setattr(engine_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None)) + monkeypatch.setattr(engine_mod, "setup_stage_devices", lambda *_: None) + monkeypatch.setattr(engine_mod, "_inject_kv_stage_info", lambda *_: None) + + def _fake_initialize_diffusion_stage(*args, **kwargs): + nonlocal captured_use_inline + captured_use_inline = kwargs.get("use_inline") + return diffusion_client + + monkeypatch.setattr(engine_mod, "initialize_diffusion_stage", _fake_initialize_diffusion_stage) + monkeypatch.setattr( + engine_mod, + "finalize_initialized_stages", + lambda stage_clients, _input_processor: ( + stage_clients, + [types.SimpleNamespace()], + [{"final_output_type": "image"}], + ), + ) + + try: + engine._initialize_stages(stage_init_timeout=1) + finally: + if old_env is None: + os.environ.pop(env_var, None) + else: + os.environ[env_var] = old_env + + assert captured_use_inline is True + def test_attach_llm_stage_uses_omni_input_preprocessor(monkeypatch): """Regression test for GLM-Image t2i preprocessing path. diff --git a/tests/entrypoints/openai_api/test_video_api_utils.py b/tests/entrypoints/openai_api/test_video_api_utils.py new file mode 100644 index 000000000..3868b6573 --- /dev/null +++ b/tests/entrypoints/openai_api/test_video_api_utils.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for OpenAI-compatible video API encoding helpers.""" + +import numpy as np +import pytest +import torch + +from vllm_omni.diffusion.postprocess import rife_interpolator +from vllm_omni.entrypoints.openai import video_api_utils + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _install_fake_video_mux(monkeypatch, mux_calls): + def _fake_mux_video_audio_bytes(frames, audio, fps, audio_sample_rate, video_codec_options=None): + mux_calls.append( + { + "frames": frames, + "audio": audio, + "fps": fps, + "audio_sample_rate": audio_sample_rate, + "video_codec_options": video_codec_options, + } + ) + return b"fake-video" + + monkeypatch.setattr( + "vllm_omni.diffusion.utils.media_utils.mux_video_audio_bytes", + _fake_mux_video_audio_bytes, + ) + + +def _install_fake_export_to_video(monkeypatch, export_calls): + def _fake_export_to_video(frames, output_path, fps): + export_calls.append( + { + "frames": frames, + "output_path": output_path, + "fps": fps, + } + ) + with open(output_path, "wb") as f: + f.write(b"fake-video") + + monkeypatch.setattr("diffusers.utils.export_to_video", _fake_export_to_video) + + +def test_encode_video_bytes_exports_frames_without_interpolation(monkeypatch): + export_calls = [] + _install_fake_export_to_video(monkeypatch, export_calls) + + frames = [np.full((2, 2, 3), fill_value=i / 5, dtype=np.float32) for i in range(5)] + video_bytes = video_api_utils._encode_video_bytes( + frames, + fps=8, + ) + + assert video_bytes == b"fake-video" + assert len(export_calls) == 1 + assert len(export_calls[0]["frames"]) == 5 + assert export_calls[0]["frames"][0].shape == (2, 2, 3) + assert export_calls[0]["fps"] == 8 + + +def test_encode_video_bytes_uses_array_fast_path(monkeypatch): + export_calls = [] + _install_fake_export_to_video(monkeypatch, export_calls) + monkeypatch.setattr( + video_api_utils, + "_normalize_frames", + lambda frames: (_ for _ in ()).throw(AssertionError("ndarray input should not go through frame list path")), + ) + + video = np.linspace(0.0, 1.0, num=5 * 2 * 2 * 3, dtype=np.float32).reshape(5, 2, 2, 3) + video_bytes = video_api_utils._encode_video_bytes(video, fps=12) + + assert video_bytes == b"fake-video" + assert len(export_calls) == 1 + assert len(export_calls[0]["frames"]) == 5 + assert export_calls[0]["frames"][0].flags["C_CONTIGUOUS"] + assert export_calls[0]["fps"] == 12.0 + + +def test_encode_video_bytes_without_audio_uses_diffusers_export(monkeypatch): + export_calls = [] + _install_fake_export_to_video(monkeypatch, export_calls) + monkeypatch.setattr( + "vllm_omni.diffusion.utils.media_utils.mux_video_audio_bytes", + lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("no-audio path should not use mux_video_audio_bytes")), + ) + + video = np.linspace(0.0, 1.0, num=4 * 2 * 2 * 3, dtype=np.float32).reshape(4, 2, 2, 3) + video_bytes = video_api_utils._encode_video_bytes(video, fps=10) + + assert video_bytes == b"fake-video" + assert len(export_calls) == 1 + assert export_calls[0]["fps"] == 10 + assert len(export_calls[0]["frames"]) == 4 + + +def test_rife_model_inference_runs_on_dummy_tensors(): + model = rife_interpolator.Model().eval() + img0 = torch.rand(1, 3, 32, 32) + img1 = torch.rand(1, 3, 32, 32) + + output = model.inference(img0, img1, scale=1.0) + + assert output.shape == (1, 3, 32, 32) + assert torch.isfinite(output).all() + + +def test_frame_interpolator_runs_actual_torch_tensor_path(monkeypatch): + model = rife_interpolator.Model().eval() + interpolator = rife_interpolator.FrameInterpolator() + monkeypatch.setattr(interpolator, "_ensure_model_loaded", lambda preferred_device=None: model) + + video = torch.zeros(1, 3, 2, 32, 32) + output_video, multiplier = interpolator.interpolate_tensor(video, exp=1, scale=1.0) + + assert multiplier == 2 + assert output_video.shape == (1, 3, 3, 32, 32) + assert torch.isfinite(output_video).all() + + +def test_frame_interpolator_prefers_input_tensor_device(monkeypatch): + chosen_devices = [] + model = rife_interpolator.Model().eval() + + def _fake_ensure_model_loaded(*, preferred_device=None): + chosen_devices.append(preferred_device) + return model + + interpolator = rife_interpolator.FrameInterpolator() + monkeypatch.setattr(interpolator, "_ensure_model_loaded", _fake_ensure_model_loaded) + monkeypatch.setattr(model.flownet, "to", lambda device: model.flownet) + + video = torch.zeros(1, 3, 2, 32, 32) + output_video, multiplier = interpolator.interpolate_tensor(video, exp=1, scale=1.0) + + assert chosen_devices == [video.device] + assert multiplier == 2 + assert output_video.shape == (1, 3, 3, 32, 32) diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 266900f70..9ebcb6897 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -37,12 +37,27 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] class MockVideoResult: - def __init__(self, videos, audios=None, sample_rate=None): + def __init__( + self, + videos, + audios=None, + sample_rate=None, + custom_output=None, + stage_durations=None, + peak_memory_mb=0.0, + ): self.multimodal_output = {"video": videos} if audios is not None: self.multimodal_output["audio"] = audios if sample_rate is not None: self.multimodal_output["audio_sample_rate"] = sample_rate + self._custom_output = custom_output or {} + self.stage_durations = stage_durations or {} + self.peak_memory_mb = peak_memory_mb + + @property + def custom_output(self): + return self._custom_output class FakeAsyncOmni: @@ -82,7 +97,7 @@ class BlockingVideoHandler: if self.stage_configs is None: self.stage_configs = stage_configs - async def generate_videos(self, request, reference_id, *, reference_image=None): + async def generate_video_bytes(self, request, reference_id, *, reference_image=None): self.started.set() try: await asyncio.Future() @@ -150,15 +165,81 @@ def _wait_until(predicate, timeout_s: float = 2.0, interval_s: float = 0.02): raise AssertionError("Timed out waiting for condition") +def test_async_video_generation_bypasses_base64(test_client, mocker: MockerFixture): + """Regression test: Ensure async video generation saves raw bytes directly + without bouncing through base64 encoding.""" + # We mock _encode_video_bytes (the correct path) + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"raw-mp4-bytes", + ) + + # We assert that encode_video_base64 is never called + mock_base64 = mocker.patch( + "vllm_omni.entrypoints.openai.serving_video.encode_video_base64", + side_effect=RuntimeError("Regression: async video path should not base64 encode"), + ) + + response = test_client.post( + "/v1/videos", + data={"prompt": "A base64 test."}, + ) + assert response.status_code == 200 + video_id = response.json()["id"] + + # Wait for completion. If it used base64, the RuntimeError would fail the task + _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value) + mock_base64.assert_not_called() + + +def test_async_video_generation_with_audio_bypasses_base64(test_client, mocker: MockerFixture): + """Regression test: Ensure async video generation passes audio through + generate_video_bytes without bouncing through base64 encoding.""" + mock_encode = mocker.patch( + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"raw-mp4-bytes", + ) + + mock_base64 = mocker.patch( + "vllm_omni.entrypoints.openai.serving_video.encode_video_base64", + side_effect=RuntimeError("Regression: async video path should not base64 encode"), + ) + + engine = test_client.app.state.openai_serving_video._engine_client + + async def _generate(prompt, request_id, sampling_params_list): + engine.captured_prompt = prompt + engine.captured_sampling_params_list = sampling_params_list + yield MockVideoResult([object()], audios=[object()], sample_rate=48000) + + engine.generate = _generate + + response = test_client.post( + "/v1/videos", + data={"prompt": "A base64 test with audio."}, + ) + assert response.status_code == 200 + video_id = response.json()["id"] + + _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value) + mock_base64.assert_not_called() + + mock_encode.assert_called_once() + kwargs = mock_encode.call_args.kwargs + assert "audio" in kwargs + assert kwargs["audio"] is not None + assert kwargs["audio_sample_rate"] == 48000 + + def test_t2v_video_generation_form(test_client, mocker: MockerFixture): fps_values = [] - def _fake_encode(video, fps): + def _fake_encode(video, fps, audio=None, audio_sample_rate=None, **kwargs): fps_values.append(fps) - return "Zg==" + return b"fake-video" mocker.patch( - "vllm_omni.entrypoints.openai.serving_video.encode_video_base64", + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", side_effect=_fake_encode, ) response = test_client.post( @@ -190,8 +271,8 @@ def test_i2v_video_generation_form(test_client, mocker: MockerFixture): image_bytes = _make_test_image_bytes((48, 32)) mocker.patch( - "vllm_omni.entrypoints.openai.serving_video.encode_video_base64", - return_value="Zg==", + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"fake-video", ) response = test_client.post( "/v1/videos", @@ -216,8 +297,8 @@ def test_i2v_video_generation_resizes_input_to_requested_dimensions(test_client, image_bytes = _make_test_image_bytes((48, 32)) mocker.patch( - "vllm_omni.entrypoints.openai.serving_video.encode_video_base64", - return_value="Zg==", + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"fake-video", ) response = test_client.post( "/v1/videos", @@ -242,8 +323,8 @@ def test_i2v_video_generation_resizes_input_to_requested_dimensions(test_client, def test_i2v_video_generation_with_image_reference_form(test_client, mocker: MockerFixture): mocker.patch( - "vllm_omni.entrypoints.openai.serving_video.encode_video_base64", - return_value="Zg==", + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"fake-video", ) response = test_client.post( "/v1/videos", @@ -267,12 +348,12 @@ def test_i2v_video_generation_with_image_reference_form(test_client, mocker: Moc def test_seconds_defaults_fps_and_frames(test_client, mocker: MockerFixture): fps_values = [] - def _fake_encode(video, fps): + def _fake_encode(video, fps, audio=None, audio_sample_rate=None, **kwargs): fps_values.append(fps) - return "Zg==" + return b"fake-video" mocker.patch( - "vllm_omni.entrypoints.openai.serving_video.encode_video_base64", + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", side_effect=_fake_encode, ) response = test_client.post( @@ -296,8 +377,8 @@ def test_seconds_defaults_fps_and_frames(test_client, mocker: MockerFixture): def test_size_param_sets_width_height(test_client, mocker: MockerFixture): mocker.patch( - "vllm_omni.entrypoints.openai.serving_video.encode_video_base64", - return_value="Zg==", + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"fake-video", ) response = test_client.post( "/v1/videos", @@ -318,8 +399,8 @@ def test_size_param_sets_width_height(test_client, mocker: MockerFixture): def test_sampling_params_pass_through(test_client, mocker: MockerFixture): mocker.patch( - "vllm_omni.entrypoints.openai.serving_video.encode_video_base64", - return_value="Zg==", + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"fake-video", ) response = test_client.post( "/v1/videos", @@ -347,13 +428,74 @@ def test_sampling_params_pass_through(test_client, mocker: MockerFixture): assert captured.extra_args["flow_shift"] == 0.25 +def test_frame_interpolation_params_pass_to_diffusion_sampling_params(test_client, mocker: MockerFixture): + """Frame interpolation parameters should be forwarded to diffusion worker sampling params.""" + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"fake-video", + ) + response = test_client.post( + "/v1/videos", + data={ + "prompt": "smooth motion", + "fps": "8", + "enable_frame_interpolation": "true", + "frame_interpolation_exp": "2", + "frame_interpolation_scale": "0.5", + "frame_interpolation_model_path": "local-rife", + }, + ) + + assert response.status_code == 200 + video_id = response.json()["id"] + _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value) + + engine = test_client.app.state.openai_serving_video._engine_client + captured = engine.captured_sampling_params_list[0] + assert captured.enable_frame_interpolation is True + assert captured.frame_interpolation_exp == 2 + assert captured.frame_interpolation_scale == 0.5 + assert captured.frame_interpolation_model_path == "local-rife" + + +def test_worker_fps_multiplier_is_applied_to_async_encoding(test_client, mocker: MockerFixture): + fps_values = [] + engine = test_client.app.state.openai_serving_video._engine_client + + async def _generate(prompt, request_id, sampling_params_list): + engine.captured_prompt = prompt + engine.captured_sampling_params_list = sampling_params_list + import numpy as np + + yield MockVideoResult([np.zeros((1, 64, 64, 3), dtype=np.uint8)], custom_output={"video_fps_multiplier": 2}) + + engine.generate = _generate + + def _fake_encode(video, fps, **kwargs): + del video, kwargs + fps_values.append(fps) + return b"fake-video" + + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + side_effect=_fake_encode, + ) + + response = test_client.post("/v1/videos", data={"prompt": "fps multiplier", "fps": "8"}) + + assert response.status_code == 200 + video_id = response.json()["id"] + _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value) + assert fps_values == [16] + + def test_audio_sample_rate_comes_from_model_config(test_client, mocker: MockerFixture): audio_sample_rates = [] - def _fake_encode(video, fps, audio=None, audio_sample_rate=None): - del video, fps, audio + def _fake_encode(video, fps, audio=None, audio_sample_rate=None, video_codec_options=None): + del video, fps, audio, video_codec_options audio_sample_rates.append(audio_sample_rate) - return "Zg==" + return b"fake-video" engine = test_client.app.state.openai_serving_video._engine_client engine.model_config = SimpleNamespace( @@ -367,12 +509,14 @@ def test_audio_sample_rate_comes_from_model_config(test_client, mocker: MockerFi async def _generate(prompt, request_id, sampling_params_list): engine.captured_prompt = prompt engine.captured_sampling_params_list = sampling_params_list - yield MockVideoResult([object()], audios=[object()]) + import numpy as np + + yield MockVideoResult([np.zeros((1, 64, 64, 3), dtype=np.uint8)], audios=[object()]) engine.generate = _generate mocker.patch( - "vllm_omni.entrypoints.openai.serving_video.encode_video_base64", + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", side_effect=_fake_encode, ) response = test_client.post( @@ -386,6 +530,33 @@ def test_audio_sample_rate_comes_from_model_config(test_client, mocker: MockerFi assert audio_sample_rates == [16000] +def test_video_job_persists_profiler_metadata(test_client, mocker: MockerFixture): + engine = test_client.app.state.openai_serving_video._engine_client + + async def _generate(prompt, request_id, sampling_params_list): + engine.captured_prompt = prompt + engine.captured_sampling_params_list = sampling_params_list + yield MockVideoResult( + [object()], + stage_durations={"diffuse": 2.5, "vae.decode": 0.3}, + peak_memory_mb=4096.5, + ) + + engine.generate = _generate + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"fake-video", + ) + + response = test_client.post("/v1/videos", data={"prompt": "profile me"}) + assert response.status_code == 200 + video_id = response.json()["id"] + completed = _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value) + + assert completed["stage_durations"] == {"diffuse": 2.5, "vae.decode": 0.3} + assert completed["peak_memory_mb"] == 4096.5 + + def test_missing_handler_returns_503(): app = FastAPI() app.include_router(router) @@ -443,8 +614,8 @@ def test_invalid_seconds_returns_422(test_client): def test_negative_prompt_and_seed_pass_through(test_client, mocker: MockerFixture): mocker.patch( - "vllm_omni.entrypoints.openai.serving_video.encode_video_base64", - return_value="Zg==", + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"fake-video", ) response = test_client.post( "/v1/videos", @@ -513,12 +684,16 @@ def test_video_request_validation(): with pytest.raises(ValueError): VideoGenerationRequest(prompt="test", image_reference={"file_id": "file-1", "image_url": "https://example.com"}) + with pytest.raises(ValueError): + VideoGenerationRequest(prompt="test", frame_interpolation_exp=0) + with pytest.raises(ValueError): + VideoGenerationRequest(prompt="test", frame_interpolation_scale=0) def test_list_videos_supports_order_after_and_limit(test_client, mocker: MockerFixture): mocker.patch( - "vllm_omni.entrypoints.openai.serving_video.encode_video_base64", - return_value="Zg==", + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"fake-video", ) ids = [] for i in range(3): @@ -586,8 +761,8 @@ def test_list_videos_supports_order_after_and_limit(test_client, mocker: MockerF def test_delete_completed_job_removes_file_and_metadata(test_client, mocker: MockerFixture): mocker.patch( - "vllm_omni.entrypoints.openai.serving_video.encode_video_base64", - return_value="Zg==", + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"fake-video", ) create_resp = test_client.post("/v1/videos", data={"prompt": "Delete this video"}) assert create_resp.status_code == 200 @@ -658,8 +833,8 @@ def test_video_response_file_extension_is_robust(): def test_extra_params_merged_into_extra_args(test_client, mocker: MockerFixture): """extra_params JSON object is merged into sampling_params.extra_args.""" mocker.patch( - "vllm_omni.entrypoints.openai.serving_video.encode_video_base64", - return_value="Zg==", + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"fake-video", ) extra_params = { "is_enable_stage2": True, @@ -689,8 +864,8 @@ def test_extra_params_merged_into_extra_args(test_client, mocker: MockerFixture) def test_extra_params_none_by_default(test_client, mocker: MockerFixture): """When extra_params is omitted, extra_args stays empty.""" mocker.patch( - "vllm_omni.entrypoints.openai.serving_video.encode_video_base64", - return_value="Zg==", + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"fake-video", ) response = test_client.post( "/v1/videos", @@ -730,8 +905,8 @@ def test_extra_params_invalid_json(test_client): def test_extra_params_merged_with_existing_extra_args(test_client, mocker: MockerFixture): """extra_params is merged on top of existing extra_args (e.g. flow_shift).""" mocker.patch( - "vllm_omni.entrypoints.openai.serving_video.encode_video_base64", - return_value="Zg==", + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"fake-video", ) response = test_client.post( "/v1/videos", @@ -752,6 +927,28 @@ def test_extra_params_merged_with_existing_extra_args(test_client, mocker: Mocke assert captured.extra_args["zero_steps"] == 2 +def test_sample_solver_forwarded_via_extra_params(test_client, mocker: MockerFixture): + """sample_solver can be passed through existing extra_params for Wan2.2 online serving.""" + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"fake-video", + ) + response = test_client.post( + "/v1/videos", + data={ + "prompt": "A fox running through snow.", + "extra_params": json.dumps({"sample_solver": "euler"}), + }, + ) + + assert response.status_code == 200 + video_id = response.json()["id"] + _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value) + engine = test_client.app.state.openai_serving_video._engine_client + captured = engine.captured_sampling_params_list[0] + assert captured.extra_args["sample_solver"] == "euler" + + @pytest.mark.asyncio async def test_run_generation_maps_omni_request_error_to_http_exception(): serving = OmniOpenAIServingVideo.for_diffusion( @@ -784,3 +981,5 @@ async def test_run_generation_maps_omni_request_error_to_http_exception(): "error_type": "OutOfMemoryError", "detail": {"retryable": False}, } + + diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index c04d46cf0..390231bf9 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -617,6 +617,49 @@ class OmniDiffusionConfig: def update_multimodal_support(self) -> None: self.supports_multimodal_inputs = self.model_class_name in {"QwenImageEditPlusPipeline"} + def enrich_config(self) -> None: + """Load model metadata from HuggingFace and populate config fields. + + Diffusers-style models expose ``model_index.json`` with ``_class_name``. + Non-diffusers models (e.g. Bagel, NextStep) only have ``config.json``, + so we fall back to reading that and mapping model_type manually. + """ + from vllm.transformers_utils.config import get_hf_file_to_dict + + try: + config_dict = get_hf_file_to_dict("model_index.json", self.model) + if config_dict is not None: + if self.model_class_name is None: + self.model_class_name = config_dict.get("_class_name", None) + self.update_multimodal_support() + + tf_config_dict = get_hf_file_to_dict("transformer/config.json", self.model) + self.tf_model_config = TransformerConfig.from_dict(tf_config_dict) + else: + raise FileNotFoundError("model_index.json not found") + except (AttributeError, OSError, ValueError, FileNotFoundError): + cfg = get_hf_file_to_dict("config.json", self.model) + if cfg is None: + raise ValueError(f"Could not find config.json or model_index.json for model {self.model}") + + self.tf_model_config = TransformerConfig.from_dict(cfg) + model_type = cfg.get("model_type") + architectures = cfg.get("architectures") or [] + + if model_type == "bagel" or "BagelForConditionalGeneration" in architectures: + self.model_class_name = "BagelPipeline" + self.tf_model_config = TransformerConfig() + self.update_multimodal_support() + elif model_type == "nextstep": + if self.model_class_name is None: + self.model_class_name = "NextStep11Pipeline" + self.tf_model_config = TransformerConfig() + self.update_multimodal_support() + elif architectures and len(architectures) == 1: + self.model_class_name = architectures[0] + else: + raise + @classmethod def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig": # Backwards-compatibility: older callers may use a diffusion-specific @@ -656,6 +699,7 @@ class DiffusionOutput: output: torch.Tensor | None = None trajectory_timesteps: list[torch.Tensor] | None = None trajectory_latents: torch.Tensor | None = None + trajectory_log_probs: torch.Tensor | None = None trajectory_decoded: list[torch.Tensor] | None = None error: str | None = None diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index c40101f0d..17dfe6d3a 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +import inspect import threading import time from collections.abc import Iterable @@ -69,6 +72,12 @@ class DiffusionEngine: self.post_process_func = get_diffusion_post_process_func(od_config) self.pre_process_func = get_diffusion_pre_process_func(od_config) + # Cache whether the model-specific postprocess accepts request-level + # sampling params so step() can support both legacy and extended hooks. + self._post_process_accepts_sampling_params = bool( + self.post_process_func is not None + and "sampling_params" in inspect.signature(self.post_process_func).parameters + ) executor_class = DiffusionExecutor.get_class(od_config) self.executor = executor_class(od_config) @@ -132,10 +141,24 @@ class DiffusionEngine: output_data = output_data.cpu() postprocess_start_time = time.perf_counter() - outputs = self.post_process_func(output_data) if self.post_process_func is not None else output_data + if self.post_process_func is not None: + # Some video pipelines need request-level controls during + # postprocess (for example worker-side frame interpolation). + if self._post_process_accepts_sampling_params: + outputs = self.post_process_func(output_data, sampling_params=request.sampling_params) + else: + outputs = self.post_process_func(output_data) + else: + outputs = output_data audio_payload = None + custom_output = output.custom_output or {} + model_audio_sample_rate = None + model_fps = None if isinstance(outputs, dict): audio_payload = outputs.get("audio") + custom_output.update(outputs.get("custom_output") or {}) + model_audio_sample_rate = outputs.get("audio_sample_rate") + model_fps = outputs.get("fps") outputs = outputs.get("video", outputs) postprocess_time = time.perf_counter() - postprocess_start_time logger.info(f"Post-processing completed in {postprocess_time:.4f} seconds") @@ -191,6 +214,10 @@ class DiffusionEngine: mm_output = {} if audio_payload is not None: mm_output["audio"] = audio_payload + if model_audio_sample_rate is not None: + mm_output["audio_sample_rate"] = model_audio_sample_rate + if model_fps is not None: + mm_output["fps"] = model_fps return [ OmniRequestOutput.from_diffusion( request_id=request_id, @@ -198,7 +225,11 @@ class DiffusionEngine: prompt=prompt, metrics=metrics, latents=output.trajectory_latents, - custom_output=output.custom_output or {}, + trajectory_latents=output.trajectory_latents, + trajectory_timesteps=output.trajectory_timesteps, + trajectory_log_probs=output.trajectory_log_probs, + trajectory_decoded=output.trajectory_decoded, + custom_output=custom_output, multimodal_output=mm_output, stage_durations=output.stage_durations, peak_memory_mb=output.peak_memory_mb, @@ -249,6 +280,10 @@ class DiffusionEngine: if num_outputs == 1: sliced_audio = sliced_audio[0] mm_output["audio"] = sliced_audio + if model_audio_sample_rate is not None: + mm_output["audio_sample_rate"] = model_audio_sample_rate + if model_fps is not None: + mm_output["fps"] = model_fps results.append( OmniRequestOutput.from_diffusion( request_id=request_id, @@ -256,7 +291,11 @@ class DiffusionEngine: prompt=prompt, metrics=metrics, latents=output.trajectory_latents, - custom_output=output.custom_output or {}, + trajectory_latents=output.trajectory_latents, + trajectory_timesteps=output.trajectory_timesteps, + trajectory_log_probs=output.trajectory_log_probs, + trajectory_decoded=output.trajectory_decoded, + custom_output=custom_output, multimodal_output=mm_output, stage_durations=output.stage_durations, peak_memory_mb=output.peak_memory_mb, diff --git a/vllm_omni/diffusion/inline_stage_diffusion_client.py b/vllm_omni/diffusion/inline_stage_diffusion_client.py new file mode 100644 index 000000000..dbf3e6a36 --- /dev/null +++ b/vllm_omni/diffusion/inline_stage_diffusion_client.py @@ -0,0 +1,362 @@ +"""Inline Stage Diffusion Client for vLLM-Omni multi-stage runtime. + +Runs DiffusionEngine in a ThreadPoolExecutor inside the Orchestrator process +instead of spawning a separate StageDiffusionProc subprocess, eliminating ZMQ +IPC overhead. Used when there is only a single diffusion stage. +""" + +from __future__ import annotations + +import asyncio +import time +from concurrent.futures import ThreadPoolExecutor +from typing import TYPE_CHECKING, Any + +import torch +from PIL import Image +from vllm.logger import init_logger + +from vllm_omni.diffusion.data import normalize_omni_error +from vllm_omni.diffusion.diffusion_engine import DiffusionEngine +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.engine.stage_init_utils import StageMetadata +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.outputs import OmniRequestOutput + +if TYPE_CHECKING: + from vllm_omni.diffusion.data import OmniDiffusionConfig + from vllm_omni.inputs.data import OmniPromptType + +logger = init_logger(__name__) + + +class InlineStageDiffusionClient: + """Runs DiffusionEngine in a thread executor inside the Orchestrator.""" + + stage_type: str = "diffusion" + + def __init__( + self, + model: str, + od_config: OmniDiffusionConfig, + metadata: StageMetadata, + batch_size: int = 1, + ) -> None: + self.model = model + self.od_config = od_config + self.stage_id = metadata.stage_id + self.final_output = metadata.final_output + self.final_output_type = metadata.final_output_type + self.default_sampling_params = metadata.default_sampling_params + self.custom_process_input_func = metadata.custom_process_input_func + self.engine_input_source = metadata.engine_input_source + self.batch_size = batch_size + + self._enrich_config() + self._engine = DiffusionEngine.make_engine(self.od_config) + self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="inline-diffusion") + + self._output_queue: asyncio.Queue[OmniRequestOutput | dict[str, Any]] = asyncio.Queue() + self._tasks: dict[str, asyncio.Task[Any]] = {} + self._shutting_down = False + + logger.info( + "[InlineStageDiffusionClient] Stage-%s initialized inline (batch_size=%d)", + self.stage_id, + self.batch_size, + ) + + def _enrich_config(self) -> None: + """Load model metadata from HuggingFace and populate od_config fields.""" + self.od_config.enrich_config() + + def _queue_error(self, request_id: str, exc: Exception, log_prefix: str) -> None: + err = normalize_omni_error(exc, request_id=request_id, stage_id=self.stage_id) + logger.exception( + "[InlineStageDiffusionClient] Stage-%s %s req=%s failed: %s", + self.stage_id, + log_prefix, + request_id, + err, + ) + self._output_queue.put_nowait( + { + "type": "error", + "request_id": request_id, + "stage_id": self.stage_id, + "status_code": err.status_code, + "error": str(err), + "error_type": err.error_type, + "detail": err.detail, + "finished": True, + } + ) + + # ------------------------------------------------------------------ + # Request processing + # ------------------------------------------------------------------ + + async def add_request_async( + self, + request_id: str, + prompt: OmniPromptType, + sampling_params: OmniDiffusionSamplingParams, + kv_sender_info: dict[int, dict[str, Any]] | None = None, + ) -> None: + task = asyncio.create_task( + self._dispatch_request( + request_id, + prompt, + sampling_params, + kv_sender_info, + ) + ) + self._tasks[request_id] = task + + async def _dispatch_request( + self, + request_id: str, + prompt: Any, + sampling_params: OmniDiffusionSamplingParams, + kv_sender_info: dict[str, Any] | None = None, + ) -> None: + try: + # Inline diffusion uses the same request envelope as AsyncOmniDiffusion. + # kv_sender_info is accepted at the interface layer for compatibility + # with other stage clients but is not consumed by DiffusionEngine here. + _ = kv_sender_info + request = OmniDiffusionRequest( + prompts=[prompt], + sampling_params=sampling_params, + request_ids=[request_id], + ) + + loop = asyncio.get_running_loop() + results = await loop.run_in_executor(self._executor, self._engine.step, request) + result = results[0] + if not result.request_id: + result.request_id = request_id + + self._output_queue.put_nowait(result) + except asyncio.CancelledError: + logger.info("request_id: %s cancelled", request_id) + except Exception as e: + self._queue_error(request_id, e, "request") + finally: + self._tasks.pop(request_id, None) + + async def add_batch_request_async( + self, + request_id: str, + prompts: list[OmniPromptType], + sampling_params: OmniDiffusionSamplingParams, + kv_sender_info: dict[int, dict[str, Any]] | None = None, + ) -> None: + task = asyncio.create_task( + self._dispatch_batch( + request_id, + prompts, + sampling_params, + kv_sender_info, + ) + ) + self._tasks[request_id] = task + + async def _dispatch_batch( + self, + request_id: str, + prompts: list[Any], + sampling_params: OmniDiffusionSamplingParams, + kv_sender_info: dict[str, Any] | None = None, + ) -> None: + try: + _ = kv_sender_info + request = OmniDiffusionRequest( + prompts=prompts, + sampling_params=sampling_params, + request_ids=[f"{request_id}-{i}" for i in range(len(prompts))], + ) + + loop = asyncio.get_running_loop() + results = await loop.run_in_executor(self._executor, self._engine.step, request) + + all_images: list = [] + merged_mm: dict[str, Any] = {} + merged_metrics: dict[str, Any] = {} + merged_durations: dict[str, float] = {} + merged_custom: dict[str, Any] = {} + peak_mem = 0.0 + latents = None + trajectory_latents: torch.Tensor | None = None + trajectory_timesteps: list[torch.Tensor] | None = None + trajectory_log_probs: torch.Tensor | None = None + trajectory_decoded: list[torch.Tensor] | list[Image.Image] | None = None + final_output_type = "image" + + for r in results: + all_images.extend(r.images) + merged_mm.update(r._multimodal_output) + merged_metrics.update(r.metrics) + merged_durations.update(r.stage_durations) + merged_custom.update(r._custom_output) + peak_mem = max(peak_mem, r.peak_memory_mb) + if latents is None and r.latents is not None: + latents = r.latents + if trajectory_latents is None: + trajectory_latents = r.trajectory_latents + if trajectory_timesteps is None: + trajectory_timesteps = r.trajectory_timesteps + if trajectory_log_probs is None: + trajectory_log_probs = r.trajectory_log_probs + if trajectory_decoded is None: + trajectory_decoded = r.trajectory_decoded + if r.final_output_type != "image": + final_output_type = r.final_output_type + + result = OmniRequestOutput.from_diffusion( + request_id=request_id, + images=all_images, + prompt=prompts[0] if len(prompts) == 1 else None, + metrics=merged_metrics, + latents=latents, + trajectory_latents=trajectory_latents, + trajectory_timesteps=trajectory_timesteps, + trajectory_log_probs=trajectory_log_probs, + trajectory_decoded=trajectory_decoded, + custom_output=merged_custom or None, + multimodal_output=merged_mm or None, + final_output_type=final_output_type, + stage_durations=merged_durations, + peak_memory_mb=peak_mem, + ) + + self._output_queue.put_nowait(result) + except asyncio.CancelledError: + logger.info("request_id: %s cancelled", request_id) + except Exception as e: + self._queue_error(request_id, e, "batch request") + finally: + self._tasks.pop(request_id, None) + + def get_diffusion_output_async(self) -> OmniRequestOutput | dict[str, Any] | None: + try: + return self._output_queue.get_nowait() + except asyncio.QueueEmpty: + return None + + def get_diffusion_output_nowait(self) -> OmniRequestOutput | dict[str, Any] | None: + return self.get_diffusion_output_async() + + async def abort_requests_async(self, request_ids: list[str]) -> None: + for rid in request_ids: + task = self._tasks.pop(rid, None) + if task: + task.cancel() + self._engine.abort(rid) + + async def collective_rpc_async( + self, + method: str, + timeout: float | None = None, + args: tuple[Any, ...] = (), + kwargs: dict[str, Any] | None = None, + ) -> Any: + loop = asyncio.get_running_loop() + + if method == "profile": + is_start = args[0] if args else True + profile_prefix = args[1] if len(args) > 1 else None + if is_start and profile_prefix is None: + profile_prefix = f"stage_{self.stage_id}_diffusion_{int(time.time())}" + return await loop.run_in_executor( + self._executor, + self._engine.profile, + is_start, + profile_prefix, + ) + + kwargs = kwargs or {} + + # LoRA methods + if method == "add_lora": + lora_request = args[0] if args else kwargs.get("lora_request") + results = await loop.run_in_executor( + self._executor, + self._engine.collective_rpc, + "add_lora", + timeout, + (), + {"lora_request": lora_request}, + None, + ) + return all(results) if isinstance(results, list) else results + + if method == "remove_lora": + results = await loop.run_in_executor( + self._executor, + self._engine.collective_rpc, + "remove_lora", + timeout, + args, + kwargs, + None, + ) + return all(results) if isinstance(results, list) else results + + if method == "list_loras": + results = await loop.run_in_executor( + self._executor, + self._engine.collective_rpc, + "list_loras", + timeout, + (), + {}, + None, + ) + if not isinstance(results, list): + return results or [] + merged: set[int] = set() + for part in results: + merged.update(part or []) + return sorted(merged) + + if method == "pin_lora": + lora_id = args[0] if args else kwargs.get("adapter_id") + results = await loop.run_in_executor( + self._executor, + self._engine.collective_rpc, + "pin_lora", + timeout, + (), + {"adapter_id": lora_id}, + None, + ) + return all(results) if isinstance(results, list) else results + + return await loop.run_in_executor( + self._executor, + self._engine.collective_rpc, + method, + timeout, + args, + kwargs, + None, + ) + + def shutdown(self) -> None: + self._shutting_down = True + + # Cancel all pending tasks + for task in self._tasks.values(): + task.cancel() + + try: + # Cancel queued futures and wait for the running one to complete deterministically + self._executor.shutdown(wait=True, cancel_futures=True) + except Exception: + pass + + try: + self._engine.close() + except Exception: + pass diff --git a/vllm_omni/diffusion/ipc.py b/vllm_omni/diffusion/ipc.py index d3d7b3aff..4f7b68db4 100644 --- a/vllm_omni/diffusion/ipc.py +++ b/vllm_omni/diffusion/ipc.py @@ -31,7 +31,13 @@ def _tensor_to_shm(tensor: torch.Tensor) -> dict[str, Any]: import numpy as np tensor = tensor.detach().cpu().contiguous() - arr = tensor.numpy() + if tensor.dtype == torch.bfloat16: + # NumPy cannot materialize BF16 tensors directly, so store the raw + # 16-bit payload and reconstruct the original dtype on unpack. + tensor_for_storage = tensor.view(torch.uint16) + arr = tensor_for_storage.numpy() + else: + arr = tensor.numpy() nbytes = arr.nbytes shm = shared_memory.SharedMemory(create=True, size=nbytes) shm_arr = np.ndarray(arr.shape, dtype=arr.dtype, buffer=shm.buf[:nbytes]) @@ -59,6 +65,8 @@ def _tensor_from_shm(handle: dict[str, Any]) -> torch.Tensor: np_dtype = np.dtype(handle["numpy_dtype"]) arr = np.ndarray(handle["shape"], dtype=np_dtype, buffer=shm.buf[: handle["nbytes"]]) tensor = torch.from_numpy(arr.copy()) + if handle["torch_dtype"] == str(torch.bfloat16): + tensor = tensor.view(torch.bfloat16) finally: shm.close() shm.unlink() @@ -72,6 +80,9 @@ def _pack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput: if output.trajectory_latents is not None and isinstance(output.trajectory_latents, torch.Tensor): if output.trajectory_latents.nelement() * output.trajectory_latents.element_size() > _SHM_TENSOR_THRESHOLD: output.trajectory_latents = _tensor_to_shm(output.trajectory_latents) + if output.trajectory_log_probs is not None and isinstance(output.trajectory_log_probs, torch.Tensor): + if output.trajectory_log_probs.nelement() * output.trajectory_log_probs.element_size() > _SHM_TENSOR_THRESHOLD: + output.trajectory_log_probs = _tensor_to_shm(output.trajectory_log_probs) return output @@ -95,6 +106,8 @@ def _unpack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput: output.output = _tensor_from_shm(output.output) if isinstance(output.trajectory_latents, dict) and output.trajectory_latents.get("__tensor_shm__"): output.trajectory_latents = _tensor_from_shm(output.trajectory_latents) + if isinstance(output.trajectory_log_probs, dict) and output.trajectory_log_probs.get("__tensor_shm__"): + output.trajectory_log_probs = _tensor_from_shm(output.trajectory_log_probs) return output diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index b780590d3..4b8d01dfa 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -25,6 +25,7 @@ from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineL from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel +from vllm_omni.diffusion.postprocess import interpolate_video_tensor from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.inputs.data import OmniTextPrompt @@ -121,10 +122,23 @@ def get_wan22_post_process_func( def post_process_func( video: torch.Tensor, output_type: str = "np", + sampling_params=None, ): if output_type == "latent": return video - return video_processor.postprocess_video(video, output_type=output_type) + custom_output = {} + if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False): + video, multiplier = interpolate_video_tensor( + video, + exp=sampling_params.frame_interpolation_exp, + scale=sampling_params.frame_interpolation_scale, + model_path=sampling_params.frame_interpolation_model_path, + ) + custom_output["video_fps_multiplier"] = multiplier + return { + "video": video_processor.postprocess_video(video, output_type=output_type), + "custom_output": custom_output, + } return post_process_func diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py index c05ecc9c9..112dfdfdd 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py @@ -30,6 +30,7 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( load_transformer_config, retrieve_latents, ) +from vllm_omni.diffusion.postprocess import interpolate_video_tensor from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.inputs.data import OmniTextPrompt @@ -72,10 +73,23 @@ def get_wan22_i2v_post_process_func( def post_process_func( video: torch.Tensor, output_type: str = "np", + sampling_params=None, ): if output_type == "latent": return video - return video_processor.postprocess_video(video, output_type=output_type) + custom_output = {} + if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False): + video, multiplier = interpolate_video_tensor( + video, + exp=sampling_params.frame_interpolation_exp, + scale=sampling_params.frame_interpolation_scale, + model_path=sampling_params.frame_interpolation_model_path, + ) + custom_output["video_fps_multiplier"] = multiplier + return { + "video": video_processor.postprocess_video(video, output_type=output_type), + "custom_output": custom_output, + } return post_process_func diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py index 261f62fb7..4953a809d 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py @@ -42,6 +42,7 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( load_transformer_config, retrieve_latents, ) +from vllm_omni.diffusion.postprocess import interpolate_video_tensor from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.platforms import current_omni_platform @@ -59,10 +60,23 @@ def get_wan22_ti2v_post_process_func( def post_process_func( video: torch.Tensor, output_type: str = "np", + sampling_params=None, ): if output_type == "latent": return video - return video_processor.postprocess_video(video, output_type=output_type) + custom_output = {} + if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False): + video, multiplier = interpolate_video_tensor( + video, + exp=sampling_params.frame_interpolation_exp, + scale=sampling_params.frame_interpolation_scale, + model_path=sampling_params.frame_interpolation_model_path, + ) + custom_output["video_fps_multiplier"] = multiplier + return { + "video": video_processor.postprocess_video(video, output_type=output_type), + "custom_output": custom_output, + } return post_process_func diff --git a/vllm_omni/diffusion/postprocess/__init__.py b/vllm_omni/diffusion/postprocess/__init__.py new file mode 100644 index 000000000..e6fe5b2d2 --- /dev/null +++ b/vllm_omni/diffusion/postprocess/__init__.py @@ -0,0 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Diffusion post-processing helpers.""" + +from vllm_omni.diffusion.postprocess.rife_interpolator import ( + FrameInterpolator, + interpolate_video_tensor, +) + +__all__ = ["FrameInterpolator", "interpolate_video_tensor"] diff --git a/vllm_omni/diffusion/postprocess/rife_interpolator.py b/vllm_omni/diffusion/postprocess/rife_interpolator.py new file mode 100644 index 000000000..b2b4a9319 --- /dev/null +++ b/vllm_omni/diffusion/postprocess/rife_interpolator.py @@ -0,0 +1,440 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +RIFE 4.22.lite frame interpolation for vLLM-Omni video generation. + +RIFE model code is vendored and adapted from: + - https://github.com/hzwer/ECCV2022-RIFE (MIT License) + - https://github.com/hzwer/Practical-RIFE (MIT License) + Copyright (c) 2021 Zhewei Huang + +The FrameInterpolator wrapper and vLLM-Omni integration code are original work. +""" + +from __future__ import annotations + +import os +import threading +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_DEFAULT_RIFE_HF_REPO = "elfgum/RIFE-4.22.lite" +_MODEL_CACHE: dict[tuple[str, str], Model] = {} +_MODEL_CACHE_LOCK = threading.Lock() + + +def warp(ten_input: torch.Tensor, ten_flow: torch.Tensor) -> torch.Tensor: + """Warp input tensor by optical flow using grid_sample.""" + ten_horizontal = ( + torch.linspace(-1.0, 1.0, ten_flow.shape[3], device=ten_flow.device) + .view(1, 1, 1, ten_flow.shape[3]) + .expand(ten_flow.shape[0], -1, ten_flow.shape[2], -1) + ) + ten_vertical = ( + torch.linspace(-1.0, 1.0, ten_flow.shape[2], device=ten_flow.device) + .view(1, 1, ten_flow.shape[2], 1) + .expand(ten_flow.shape[0], -1, -1, ten_flow.shape[3]) + ) + ten_grid = torch.cat([ten_horizontal, ten_vertical], dim=1) + + ten_flow = torch.cat( + [ + ten_flow[:, 0:1, :, :] / ((ten_input.shape[3] - 1.0) / 2.0), + ten_flow[:, 1:2, :, :] / ((ten_input.shape[2] - 1.0) / 2.0), + ], + dim=1, + ) + grid = (ten_grid + ten_flow).permute(0, 2, 3, 1) + return F.grid_sample( + input=ten_input, + grid=grid, + mode="bilinear", + padding_mode="border", + align_corners=True, + ) + + +def _conv( + in_planes: int, + out_planes: int, + kernel_size: int = 3, + stride: int = 1, + padding: int = 1, + dilation: int = 1, +) -> nn.Sequential: + return nn.Sequential( + nn.Conv2d( + in_planes, + out_planes, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=True, + ), + nn.LeakyReLU(0.2, True), + ) + + +class ResConv(nn.Module): + """Residual convolution block with learnable beta scaling.""" + + def __init__(self, c: int, dilation: int = 1): + super().__init__() + self.conv = nn.Conv2d(c, c, 3, 1, dilation, dilation=dilation, groups=1) + self.beta = nn.Parameter(torch.ones((1, c, 1, 1)), requires_grad=True) + self.relu = nn.LeakyReLU(0.2, True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.relu(self.conv(x) * self.beta + x) + + +class IFBlock(nn.Module): + """Single-scale optical flow, mask, and feature block.""" + + def __init__(self, in_planes: int, c: int = 64): + super().__init__() + self.conv0 = nn.Sequential( + _conv(in_planes, c // 2, 3, 2, 1), + _conv(c // 2, c, 3, 2, 1), + ) + self.convblock = nn.Sequential( + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ) + self.lastconv = nn.Sequential( + nn.ConvTranspose2d(c, 4 * 13, 4, 2, 1), + nn.PixelShuffle(2), + ) + + def forward( + self, + x: torch.Tensor, + flow: torch.Tensor | None = None, + scale: float = 1.0, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + x = F.interpolate(x, scale_factor=1.0 / scale, mode="bilinear", align_corners=False) + if flow is not None: + flow = ( + F.interpolate( + flow, + scale_factor=1.0 / scale, + mode="bilinear", + align_corners=False, + ) + * 1.0 + / scale + ) + x = torch.cat((x, flow), 1) + feat = self.conv0(x) + feat = self.convblock(feat) + tmp = self.lastconv(feat) + tmp = F.interpolate(tmp, scale_factor=scale, mode="bilinear", align_corners=False) + flow = tmp[:, :4] * scale + mask = tmp[:, 4:5] + feat = tmp[:, 5:] + return flow, mask, feat + + +class Head(nn.Module): + """Feature encoder producing four-channel features at full resolution.""" + + def __init__(self): + super().__init__() + self.cnn0 = nn.Conv2d(3, 16, 3, 2, 1) + self.cnn1 = nn.Conv2d(16, 16, 3, 1, 1) + self.cnn2 = nn.Conv2d(16, 16, 3, 1, 1) + self.cnn3 = nn.ConvTranspose2d(16, 4, 4, 2, 1) + self.relu = nn.LeakyReLU(0.2, True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x0 = self.cnn0(x) + x = self.relu(x0) + x1 = self.cnn1(x) + x = self.relu(x1) + x2 = self.cnn2(x) + x = self.relu(x2) + x3 = self.cnn3(x) + return x3 + + +class IFNet(nn.Module): + """Four-scale IFNet optical flow network.""" + + def __init__(self): + super().__init__() + self.block0 = IFBlock(7 + 8, c=192) + self.block1 = IFBlock(8 + 4 + 8 + 8, c=128) + self.block2 = IFBlock(8 + 4 + 8 + 8, c=64) + self.block3 = IFBlock(8 + 4 + 8 + 8, c=32) + self.encode = Head() + + def forward( + self, + x: torch.Tensor, + timestep: float = 0.5, + scale_list: list[float] | None = None, + ) -> tuple[list[torch.Tensor], torch.Tensor, list[tuple[torch.Tensor, torch.Tensor] | torch.Tensor]]: + if scale_list is None: + scale_list = [8, 4, 2, 1] + + channel = x.shape[1] // 2 + img0 = x[:, :channel] + img1 = x[:, channel:] + + if not torch.is_tensor(timestep): + timestep = (x[:, :1].clone() * 0 + 1) * timestep + else: + timestep = timestep.repeat(1, 1, img0.shape[2], img0.shape[3]) + + f0 = self.encode(img0[:, :3]) + f1 = self.encode(img1[:, :3]) + + flow_list: list[torch.Tensor] = [] + merged: list[tuple[torch.Tensor, torch.Tensor] | torch.Tensor] = [] + mask_list: list[torch.Tensor] = [] + warped_img0 = img0 + warped_img1 = img1 + flow = None + mask = None + + for i, block in enumerate([self.block0, self.block1, self.block2, self.block3]): + if flow is None: + flow, mask, feat = block( + torch.cat((img0[:, :3], img1[:, :3], f0, f1, timestep), 1), + None, + scale=scale_list[i], + ) + else: + wf0 = warp(f0, flow[:, :2]) + wf1 = warp(f1, flow[:, 2:4]) + fd, m0, feat = block( + torch.cat( + ( + warped_img0[:, :3], + warped_img1[:, :3], + wf0, + wf1, + timestep, + mask, + feat, + ), + 1, + ), + flow, + scale=scale_list[i], + ) + mask = m0 + flow = flow + fd + + mask_list.append(mask) + flow_list.append(flow) + warped_img0 = warp(img0, flow[:, :2]) + warped_img1 = warp(img1, flow[:, 2:4]) + merged.append((warped_img0, warped_img1)) + + mask = torch.sigmoid(mask) + merged[3] = warped_img0 * mask + warped_img1 * (1 - mask) + return flow_list, mask_list[3], merged + + +class Model: + """Wraps IFNet and exposes RIFE-compatible load/inference helpers.""" + + def __init__(self): + self.flownet = IFNet() + + def eval(self) -> Model: + self.flownet.eval() + return self + + def device(self) -> torch.device: + return next(self.flownet.parameters()).device + + def load_model(self, path: str) -> None: + flownet_path = os.path.join(path, "flownet.pkl") + if not os.path.isfile(flownet_path): + raise FileNotFoundError( + f"RIFE weight file not found: {flownet_path}. Expected layout: /flownet.pkl" + ) + + state = torch.load(flownet_path, map_location="cpu", weights_only=False) + state = {k.removeprefix("module."): v for k, v in state.items()} + self.flownet.load_state_dict(state, strict=False) + logger.info("Loaded RIFE weights from %s", flownet_path) + + def inference( + self, + img0: torch.Tensor, + img1: torch.Tensor, + scale: float = 1.0, + timestep: float = 0.5, + ) -> torch.Tensor: + _n, _c, h, w = img0.shape + ph = ((h - 1) // 32 + 1) * 32 + pw = ((w - 1) // 32 + 1) * 32 + pad = (0, pw - w, 0, ph - h) + img0 = F.pad(img0, pad) + img1 = F.pad(img1, pad) + + imgs = torch.cat((img0, img1), 1) + scale_list = [8 / scale, 4 / scale, 2 / scale, 1 / scale] + with torch.no_grad(): + _flow_list, _mask, merged = self.flownet( + imgs, + timestep=timestep, + scale_list=scale_list, + ) + return merged[3][:, :, :h, :w] + + +def _resolve_rife_model_path(model_path: str | None) -> str: + model_path = model_path or _DEFAULT_RIFE_HF_REPO + if os.path.isdir(model_path): + return model_path + from vllm_omni.model_executor.model_loader.weight_utils import ( + download_weights_from_hf_specific, + ) + + return download_weights_from_hf_specific( + model_path, + cache_dir=None, + allow_patterns=["flownet.pkl"], + require_all=True, + ) + + +def _select_torch_device() -> torch.device: + try: + from vllm_omni.platforms import current_omni_platform + + return current_omni_platform.get_torch_device() + except Exception as exc: + logger.warning("Failed to resolve current vLLM-Omni torch device: %s", exc) + + if torch.cuda.is_available(): + return torch.device("cuda") + return torch.device("cpu") + + +def _normalize_video_tensor_layout(video: torch.Tensor) -> tuple[torch.Tensor, Any]: + if video.ndim == 5: + if video.shape[1] in (3, 4): + return video, lambda out: out + if video.shape[2] in (3, 4): + return video.permute(0, 2, 1, 3, 4), lambda out: out.permute(0, 2, 1, 3, 4) + elif video.ndim == 4: + if video.shape[0] in (3, 4): + return video.unsqueeze(0), lambda out: out.squeeze(0) + if video.shape[1] in (3, 4): + return video.permute(1, 0, 2, 3).unsqueeze(0), lambda out: out.squeeze(0).permute(1, 0, 2, 3) + raise ValueError(f"Unsupported video tensor shape for interpolation: {tuple(video.shape)}") + + +def _normalize_video_tensor_range(video: torch.Tensor) -> tuple[torch.Tensor, Any]: + original_dtype = video.dtype + video = video.detach() + if video.is_floating_point(): + video = video.to(torch.float32) + if torch.amin(video) < 0.0 or torch.amax(video) > 1.0: + return video.clamp(-1.0, 1.0) * 0.5 + 0.5, lambda out: (out * 2.0 - 1.0).to(original_dtype) + return video.clamp(0.0, 1.0), lambda out: out.to(original_dtype) + return video.to(torch.float32) / 255.0, lambda out: (out * 255.0).round().clamp(0, 255).to(original_dtype) + + +class FrameInterpolator: + """Lazy-loaded RIFE 4.22.lite frame interpolator.""" + + def __init__(self, model_path: str | None = None): + self._model_path = model_path + self._resolved_path: str | None = None + + def _ensure_model_loaded(self, preferred_device: torch.device | None = None) -> Model: + resolved_path = _resolve_rife_model_path(self._model_path) + self._resolved_path = resolved_path + device = preferred_device or _select_torch_device() + cache_key = (resolved_path, str(device)) + + with _MODEL_CACHE_LOCK: + if cache_key in _MODEL_CACHE: + return _MODEL_CACHE[cache_key] + + model = Model() + model.load_model(resolved_path) + model.eval() + model.flownet = model.flownet.to(device) + _MODEL_CACHE[cache_key] = model + logger.info("RIFE model loaded on device: %s", device) + return model + + def _make_inference( + self, + model: Model, + img0: torch.Tensor, + img1: torch.Tensor, + n: int, + scale: float, + ) -> list[torch.Tensor]: + if n == 1: + return [model.inference(img0, img1, scale=scale)] + mid = model.inference(img0, img1, scale=scale) + return ( + self._make_inference(model, img0, mid, n // 2, scale) + + [mid] + + self._make_inference(model, mid, img1, n // 2, scale) + ) + + def interpolate_tensor( + self, + video: torch.Tensor, + exp: int = 1, + scale: float = 1.0, + ) -> tuple[torch.Tensor, int]: + if exp < 1: + raise ValueError(f"frame interpolation exp must be >= 1, got {exp}") + if scale <= 0: + raise ValueError(f"frame interpolation scale must be > 0, got {scale}") + + video, restore_layout = _normalize_video_tensor_layout(video) + if video.shape[2] < 2: + return restore_layout(video), 1 + + video, restore_range = _normalize_video_tensor_range(video) + # Prefer the decoded video's current device so CPU-offloaded requests do + # not move the tensor back to GPU just for interpolation. + model = self._ensure_model_loaded(preferred_device=video.device) + video = video.to(model.device()) + intermediates_per_pair = 2**exp // 2 + + result_frames: list[torch.Tensor] = [] + for idx in range(video.shape[2] - 1): + img0 = video[:, :, idx, :, :] + img1 = video[:, :, idx + 1, :, :] + result_frames.append(img0) + result_frames.extend(self._make_inference(model, img0, img1, intermediates_per_pair, scale)) + result_frames.append(video[:, :, -1, :, :]) + result = torch.stack(result_frames, dim=2) + return restore_layout(restore_range(result)), 2**exp + + +def interpolate_video_tensor( + video: torch.Tensor, + exp: int = 1, + scale: float = 1.0, + model_path: str | None = None, +) -> tuple[torch.Tensor, int]: + """Interpolate a video tensor and return the FPS multiplier.""" + interpolator = FrameInterpolator(model_path=model_path) + return interpolator.interpolate_tensor(video, exp=exp, scale=scale) diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py index f11d300f7..93aeabea1 100644 --- a/vllm_omni/diffusion/stage_diffusion_client.py +++ b/vllm_omni/diffusion/stage_diffusion_client.py @@ -24,6 +24,21 @@ if TYPE_CHECKING: logger = init_logger(__name__) +def create_diffusion_client( + model: str, + od_config: OmniDiffusionConfig, + metadata: StageMetadata, + batch_size: int = 1, + use_inline: bool = False, +) -> Any: + """Factory to create either an inline or the default diffusion client.""" + if use_inline: + from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient + + return InlineStageDiffusionClient(model, od_config, metadata, batch_size=batch_size) + return StageDiffusionClient(model, od_config, metadata, batch_size=batch_size) + + class StageDiffusionClient: """Wraps AsyncOmniDiffusion for use inside the Orchestrator. @@ -176,10 +191,8 @@ class StageDiffusionClient: "todo": True, "reason": f"AsyncOmniDiffusion.{method} is not implemented", } - # Extract is_start and profile_prefix from args is_start = args[0] if args else True profile_prefix = args[1] if len(args) > 1 else None - # Generate profile_prefix with stage_id if starting and no prefix provided if is_start and profile_prefix is None: profile_prefix = f"stage_{self.stage_id}_diffusion_{int(time.time())}" result = target(is_start, profile_prefix) @@ -200,7 +213,6 @@ class StageDiffusionClient: return await asyncio.wait_for(result, timeout=timeout) return await result - # Fall back to collective RPC for other methods loop = asyncio.get_running_loop() return await loop.run_in_executor( self._engine._executor, diff --git a/vllm_omni/diffusion/utils/media_utils.py b/vllm_omni/diffusion/utils/media_utils.py new file mode 100644 index 000000000..a09cd4595 --- /dev/null +++ b/vllm_omni/diffusion/utils/media_utils.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Video/audio muxing utilities using PyAV (no ffmpeg binary dependency).""" + +from __future__ import annotations + +import io +from fractions import Fraction + +import av +import numpy as np + + +def mux_video_audio_bytes( + video_frames: np.ndarray, + audio_waveform: np.ndarray | None = None, + *, + fps: float = 25.0, + audio_sample_rate: int = 44100, + video_codec: str = "h264", + audio_codec: str = "aac", + crf: str = "18", + video_codec_options: dict[str, str] | None = None, +) -> bytes: + """Mux video frames and optional audio waveform into MP4 bytes. + + Args: + video_frames: uint8 array of shape ``(T, H, W, 3)`` (RGB). + audio_waveform: float32 array – mono ``(N,)`` or ``(N, C)`` / ``(C, N)``. + fps: Video frame rate. + audio_sample_rate: Audio sample rate in Hz. + video_codec: Video codec name. + audio_codec: Audio codec name. + crf: Constant rate factor for the video encoder. + + Returns: + Raw MP4 bytes ready to be written to disk or streamed. + """ + buf = io.BytesIO() + container = av.open(buf, mode="w", format="mp4") + + v_stream = container.add_stream(video_codec, rate=Fraction(fps).limit_denominator(10000)) + v_stream.width = video_frames.shape[2] + v_stream.height = video_frames.shape[1] + v_stream.pix_fmt = "yuv420p" + + options = {"crf": str(crf)} + if video_codec_options: + options.update(video_codec_options) + v_stream.options = options + + a_stream = None + if audio_waveform is not None: + samples = audio_waveform.astype(np.float32) + if samples.ndim == 1: + samples = samples.reshape(1, -1) + elif samples.ndim == 2 and samples.shape[0] > samples.shape[1]: + samples = np.ascontiguousarray(samples.T) + num_channels = samples.shape[0] + layout = "stereo" if num_channels >= 2 else "mono" + a_stream = container.add_stream(audio_codec, rate=audio_sample_rate) + a_stream.layout = layout + + for frame_data in video_frames: + frame = av.VideoFrame.from_ndarray(frame_data, format="rgb24") + for packet in v_stream.encode(frame): + container.mux(packet) + for packet in v_stream.encode(): + container.mux(packet) + + if a_stream is not None and audio_waveform is not None: + audio_frame = av.AudioFrame.from_ndarray(samples, format="fltp", layout=layout) + audio_frame.sample_rate = audio_sample_rate + for packet in a_stream.encode(audio_frame): + container.mux(packet) + for packet in a_stream.encode(): + container.mux(packet) + + container.close() + return buf.getvalue() diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index a4d87c96e..34d9d2404 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -487,11 +487,13 @@ class AsyncOmniEngine: inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to) _inject_kv_stage_info(stage_cfg, stage_id) + use_inline = self.num_stages == 1 stage_clients[stage_id] = initialize_diffusion_stage( self.model, stage_cfg, metadata, batch_size=self.diffusion_batch_size, + use_inline=use_inline, ) logger.info( "[AsyncOmniEngine] Stage %s initialized (diffusion, batch_size=%d)", diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index b2ce6ec70..a370da8e1 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -267,6 +267,23 @@ class Orchestrator: req_state = self.request_states.get(output.request_id) if req_state is not None: + if getattr(output, "error", None) is not None: + parent_id = self._companion_to_parent.get(output.request_id, output.request_id) + await self.output_async_queue.put( + { + "type": "error", + "request_id": parent_id, + "stage_id": stage_id, + "error": output.error, + } + ) + role_map = self._companion_map.get(parent_id, {}) + for cid in role_map.values(): + self.request_states.pop(cid, None) + self._cleanup_companion_state(parent_id) + self.request_states.pop(parent_id, None) + continue + stage_metrics = self._build_stage_metrics(stage_id, output.request_id, [output], req_state) await self._route_output(stage_id, output, req_state, stage_metrics) continue diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 9c246ce6e..ffe90cb3b 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -437,6 +437,7 @@ def initialize_diffusion_stage( stage_cfg: Any, metadata: StageMetadata, batch_size: int = 1, + use_inline: bool = False, ) -> Any: """Build a diffusion stage client. @@ -447,9 +448,11 @@ def initialize_diffusion_stage( batch_size: Maximum number of requests to batch together in the diffusion engine. Passed through to ``StageDiffusionClient`` and ultimately to ``AsyncOmniDiffusion``. + use_inline: If True, uses the inline diffusion client instead of subprocess. """ + from vllm_omni.diffusion.stage_diffusion_client import create_diffusion_client + from vllm_omni.diffusion.data import OmniDiffusionConfig - from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient od_config = OmniDiffusionConfig.from_kwargs( model=model, @@ -457,7 +460,7 @@ def initialize_diffusion_stage( ) if metadata.cfg_kv_collect_func is not None: od_config.cfg_kv_collect_func = metadata.cfg_kv_collect_func - return StageDiffusionClient(model, od_config, metadata, batch_size=batch_size) + return create_diffusion_client(model, od_config, metadata, batch_size=batch_size, use_inline=use_inline) def close_started_llm_stage(started: StartedLlmStage) -> None: diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index ceb87f111..cec9869d2 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -1917,18 +1917,6 @@ def video_response_from_request(model_name: str, req: VideoGenerationRequest) -> return resp -async def decode_and_save_video_output(output: Any, file_name: str) -> str: - if not output.b64_json: - raise RuntimeError(f"Video output for {file_name} did not include b64_json content.") - - try: - video_bytes = base64.b64decode(output.b64_json) - except Exception as decode_exc: - raise RuntimeError(f"Failed to decode generated video payload for {file_name}") from decode_exc - - return await STORAGE_MANAGER.save(video_bytes, file_name) - - def _cleanup_video(video_id: str, output_path: str | None): try: if output_path is not None: @@ -1952,15 +1940,12 @@ async def _run_video_generation_job( started_at = time.perf_counter() output_path = None try: - response = await handler.generate_videos(request, video_id, reference_image=reference_image) - if not response.data: - raise RuntimeError("Video generation completed but returned no outputs.") - - if (video_count := len(response.data)) > 1: - logger.warning("Video request %s generated %s outputs but we only expected one.", video_id, video_count) + video_bytes, stage_durations, peak_memory_mb = await handler.generate_video_bytes( + request, video_id, reference_image=reference_image + ) file_name = f"{video_id}.{job.file_extension}" - output_path = await decode_and_save_video_output(response.data[0], file_name) + output_path = await STORAGE_MANAGER.save(video_bytes, file_name) logger.info("Video request %s persisted %s output file.", video_id, output_path) await VIDEO_STORE.update_fields( @@ -1971,6 +1956,8 @@ async def _run_video_generation_job( "file_name": file_name, "completed_at": int(time.time()), "inference_time_s": time.perf_counter() - started_at, + "stage_durations": stage_durations, + "peak_memory_mb": peak_memory_mb, }, ) except Exception as exc: @@ -2023,6 +2010,10 @@ async def create_video( true_cfg_scale: float | None = Form(default=None), seed: int | None = Form(default=None), negative_prompt: str | None = Form(default=None), + enable_frame_interpolation: bool = Form(default=False), + frame_interpolation_exp: int = Form(default=1, ge=1), + frame_interpolation_scale: float = Form(default=1.0, gt=0.0), + frame_interpolation_model_path: str | None = Form(default=None), lora: str | None = Form(default=None), extra_params: str | None = Form(default=None), ) -> VideoResponse: @@ -2092,6 +2083,10 @@ async def create_video( "true_cfg_scale": true_cfg_scale, "seed": seed, "negative_prompt": negative_prompt, + "enable_frame_interpolation": enable_frame_interpolation, + "frame_interpolation_exp": frame_interpolation_exp, + "frame_interpolation_scale": frame_interpolation_scale, + "frame_interpolation_model_path": frame_interpolation_model_path, "lora": _parse_form_json(lora, expected_type=dict), "extra_params": _parse_form_json(extra_params, expected_type=dict), } diff --git a/vllm_omni/entrypoints/openai/protocol/videos.py b/vllm_omni/entrypoints/openai/protocol/videos.py index e180bef22..d23bc1d2e 100644 --- a/vllm_omni/entrypoints/openai/protocol/videos.py +++ b/vllm_omni/entrypoints/openai/protocol/videos.py @@ -150,6 +150,29 @@ class VideoGenerationRequest(BaseModel): ) seed: int | None = Field(default=None, description="Random seed for reproducibility") + # vllm-omni extensions for post-generation frame interpolation. + enable_frame_interpolation: bool = Field( + default=False, + description="Enable post-generation RIFE frame interpolation before MP4 encoding.", + ) + frame_interpolation_exp: int = Field( + default=1, + ge=1, + description="Interpolation exponent: 1=2x temporal resolution, 2=4x, etc.", + ) + frame_interpolation_scale: float = Field( + default=1.0, + gt=0.0, + description="RIFE inference scale. Use 0.5 for high-resolution inputs to save memory.", + ) + frame_interpolation_model_path: str | None = Field( + default=None, + description=( + "Local directory or Hugging Face repo ID containing RIFE flownet.pkl weights. " + "Defaults to elfgum/RIFE-4.22.lite." + ), + ) + # vllm-omni extension for per-request LoRA. lora: dict[str, Any] | None = Field( default=None, diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index 390dd2919..50bfaf7f4 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/vllm_omni/entrypoints/openai/serving_video.py @@ -21,7 +21,7 @@ from vllm_omni.entrypoints.openai.protocol.videos import ( VideoGenerationResponse, ) from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_request -from vllm_omni.entrypoints.openai.video_api_utils import encode_video_base64 +from vllm_omni.entrypoints.openai.video_api_utils import _encode_video_bytes, encode_video_base64 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams, OmniTextPrompt logger = init_logger(__name__) @@ -29,11 +29,23 @@ logger = init_logger(__name__) @dataclass class ReferenceImage: - """Reference class for tracking additional metadata if needed""" + """Reference class for tracking additional metadata if needed.""" data: Image.Image +@dataclass +class VideoGenerationArtifacts: + """Normalized outputs and profiler metadata extracted from one request.""" + + videos: list[Any] + audios: list[Any | None] + audio_sample_rate: int + output_fps: int + stage_durations: dict[str, float] + peak_memory_mb: float + + class OmniOpenAIServingVideo: """OpenAI-style video generation handler for omni diffusion models.""" @@ -72,13 +84,14 @@ class OmniOpenAIServingVideo: stage_configs=stage_configs, ) - async def generate_videos( + async def _run_and_extract( self, request: VideoGenerationRequest, reference_id: str, *, reference_image: ReferenceImage | None = None, - ) -> VideoGenerationResponse: + ) -> VideoGenerationArtifacts: + """Run the generation pipeline and extract video/audio/profiler outputs.""" prompt: OmniTextPrompt = OmniTextPrompt(prompt=request.prompt) if request.negative_prompt is not None: prompt["negative_prompt"] = request.negative_prompt @@ -101,6 +114,10 @@ class OmniOpenAIServingVideo: if vp.fps is not None: gen_params.fps = vp.fps gen_params.frame_rate = float(vp.fps) + gen_params.enable_frame_interpolation = request.enable_frame_interpolation + gen_params.frame_interpolation_exp = request.frame_interpolation_exp + gen_params.frame_interpolation_scale = request.frame_interpolation_scale + gen_params.frame_interpolation_model_path = request.frame_interpolation_model_path if request.num_inference_steps is not None: gen_params.num_inference_steps = request.num_inference_steps @@ -123,14 +140,12 @@ class OmniOpenAIServingVideo: if request.flow_shift is not None: gen_params.extra_args["flow_shift"] = request.flow_shift - # Apply model-specific extra parameters if request.extra_params is not None: if not isinstance(request.extra_params, dict): raise HTTPException( status_code=HTTPStatus.BAD_REQUEST.value, detail="extra_params must be a JSON object/dict.", ) - # Merge extra_params into extra_args gen_params.extra_args.update(request.extra_params) logger.info("Applied extra_params: %s", request.extra_params) @@ -145,31 +160,103 @@ class OmniOpenAIServingVideo: ) result = await self._run_generation(prompt, gen_params, reference_id) - _t_encode_start = time.perf_counter() videos = self._extract_video_outputs(result) audios = self._extract_audio_outputs(result, expected_count=len(videos)) audio_sample_rate = self._resolve_audio_sample_rate(result) - output_fps = vp.fps or 24 + output_fps = (vp.fps or self._resolve_fps(result) or 24) * self._resolve_video_fps_multiplier(result) + return VideoGenerationArtifacts( + videos=videos, + audios=audios, + audio_sample_rate=audio_sample_rate, + output_fps=output_fps, + stage_durations=self._extract_stage_durations(result), + peak_memory_mb=self._extract_peak_memory_mb(result), + ) + async def generate_videos( + self, + request: VideoGenerationRequest, + reference_id: str, + *, + reference_image: ReferenceImage | None = None, + ) -> VideoGenerationResponse: + artifacts = await self._run_and_extract(request, reference_id, reference_image=reference_image) + + video_codec_options = self._resolve_video_codec_options(request) + _t_encode_start = time.perf_counter() video_data = [ VideoData( b64_json=( - encode_video_base64(video, fps=output_fps) - if audios[idx] is None + encode_video_base64(video, fps=artifacts.output_fps, video_codec_options=video_codec_options) + if artifacts.audios[idx] is None else encode_video_base64( video, - fps=output_fps, - audio=audios[idx], - audio_sample_rate=audio_sample_rate, + fps=artifacts.output_fps, + audio=artifacts.audios[idx], + audio_sample_rate=artifacts.audio_sample_rate, + video_codec_options=video_codec_options, ) ) ) - for idx, video in enumerate(videos) + for idx, video in enumerate(artifacts.videos) ] _t_encode_ms = (time.perf_counter() - _t_encode_start) * 1000 logger.info("Video response encoding (MP4+base64): %.2f ms", _t_encode_ms) return VideoGenerationResponse(created=int(time.time()), data=video_data) + async def generate_video_bytes( + self, + request: VideoGenerationRequest, + reference_id: str, + *, + reference_image: ReferenceImage | None = None, + ) -> tuple[bytes, dict[str, float], float]: + """Generate a video and return raw MP4 bytes, bypassing base64 encoding.""" + artifacts = await self._run_and_extract(request, reference_id, reference_image=reference_image) + if len(artifacts.videos) > 1: + logger.warning( + "Video request %s generated %d outputs; returning only the first.", + reference_id, + len(artifacts.videos), + ) + audio = artifacts.audios[0] + + _t_encode_start = time.perf_counter() + 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=self._resolve_video_codec_options(request), + ) + _t_encode_ms = (time.perf_counter() - _t_encode_start) * 1000 + logger.info("Video response encoding (MP4 bytes): %.2f ms", _t_encode_ms) + return video_bytes, artifacts.stage_durations, artifacts.peak_memory_mb + + @staticmethod + def _resolve_video_codec_options(request: VideoGenerationRequest) -> dict[str, str]: + video_codec_options = {"preset": "ultrafast", "threads": "0"} + if request.extra_params is not None and isinstance(request.extra_params, dict): + extra_codec_options = request.extra_params.get("video_codec_options") + if isinstance(extra_codec_options, dict): + video_codec_options = {str(k): str(v) for k, v in extra_codec_options.items()} + return video_codec_options + + @staticmethod + def _resolve_video_fps_multiplier(result: Any) -> int: + custom_output = getattr(result, "custom_output", None) + if isinstance(custom_output, dict): + multiplier = custom_output.get("video_fps_multiplier") + if multiplier is not None: + return int(multiplier) + request_output = getattr(result, "request_output", None) + if request_output is not None: + custom_output = getattr(request_output, "custom_output", None) + if isinstance(custom_output, dict): + multiplier = custom_output.get("video_fps_multiplier") + if multiplier is not None: + return int(multiplier) + return 1 + @staticmethod def _apply_lora(lora_body: Any, gen_params: OmniDiffusionSamplingParams) -> None: try: @@ -201,7 +288,6 @@ class OmniOpenAIServingVideo: detail="Stage configs not found. Start server with an omni diffusion model.", ) - # Video generation endpoint only supports diffusion stages. for stage in stage_configs: stage_type = get_stage_type(stage) if stage_type != "diffusion": @@ -210,7 +296,6 @@ class OmniOpenAIServingVideo: detail=f"Video generation only supports diffusion stages, found '{stage_type}' stage.", ) - # Common generation logic for both paths engine_client = cast(AsyncOmni, self._engine_client) sampling_params_list: list[OmniSamplingParams] = [gen_params for _ in stage_configs] @@ -336,6 +421,45 @@ class OmniOpenAIServingVideo: return 24000 + @staticmethod + def _resolve_fps(result: Any) -> int | None: + multimodal_output = getattr(result, "multimodal_output", None) + if isinstance(multimodal_output, dict): + fps = multimodal_output.get("fps") + if fps is not None: + try: + fps_val = fps.item() if hasattr(fps, "item") else int(fps) + if fps_val > 0: + return fps_val + except (TypeError, ValueError): + pass + + request_output = getattr(result, "request_output", None) + if isinstance(request_output, dict): + mm = request_output.get("multimodal_output") or {} + if isinstance(mm, dict): + fps = mm.get("fps") + if fps is not None: + try: + fps_val = fps.item() if hasattr(fps, "item") else int(fps) + if fps_val > 0: + return fps_val + except (TypeError, ValueError): + pass + elif hasattr(request_output, "multimodal_output"): + mm = getattr(request_output, "multimodal_output", None) + if isinstance(mm, dict): + fps = mm.get("fps") + if fps is not None: + try: + fps_val = fps.item() if hasattr(fps, "item") else int(fps) + if fps_val > 0: + return fps_val + except (TypeError, ValueError): + pass + + return None + @classmethod def _extract_audio_sample_rate_from_result(cls, result: Any) -> int | None: multimodal_output = getattr(result, "multimodal_output", None) @@ -414,3 +538,16 @@ class OmniOpenAIServingVideo: return None return sample_rate if sample_rate > 0 else None + + @staticmethod + def _extract_stage_durations(result: Any) -> dict[str, float]: + stage_durations = getattr(result, "stage_durations", None) + return stage_durations if isinstance(stage_durations, dict) else {} + + @staticmethod + def _extract_peak_memory_mb(result: Any) -> float: + peak_memory_mb = getattr(result, "peak_memory_mb", 0.0) + try: + return float(peak_memory_mb or 0.0) + except (TypeError, ValueError): + return 0.0 diff --git a/vllm_omni/entrypoints/openai/video_api_utils.py b/vllm_omni/entrypoints/openai/video_api_utils.py index 2ed1fd3de..d118d5149 100644 --- a/vllm_omni/entrypoints/openai/video_api_utils.py +++ b/vllm_omni/entrypoints/openai/video_api_utils.py @@ -159,111 +159,116 @@ def _normalize_frames(frames: list[Any]) -> list[np.ndarray]: return normalized -def _coerce_video_to_frames(video: Any) -> list[np.ndarray]: - """Convert a video payload into a list of frames for export_to_video.""" +def _coerce_video_to_array(video: Any) -> np.ndarray: + """Convert a video payload into a normalized array of frames (F, H, W, C).""" if isinstance(video, torch.Tensor): video_array = _normalize_video_tensor(video) - return list(video_array) + if video_array.ndim == 4: + return video_array + if video_array.ndim == 3: + return video_array[None, ...] + raise ValueError(f"Unsupported video array shape: {video_array.shape}") if isinstance(video, np.ndarray): video_array = _normalize_video_array(video) if isinstance(video_array, list): raise ValueError("Batched video arrays must be split before encoding.") if video_array.ndim == 4: - return list(video_array) + return video_array if video_array.ndim == 3: - return [video_array] + return video_array[None, ...] raise ValueError(f"Unsupported video array shape: {video_array.shape}") if isinstance(video, list): if not video: - return [] + return np.empty((0,), dtype=np.float32) # If this looks like a list of frames, normalize directly. if all(isinstance(item, (np.ndarray, torch.Tensor, Image.Image)) for item in video): # If each item is itself a video (ndim==4), handle elsewhere. if all(hasattr(item, "ndim") and item.ndim >= 4 for item in video): raise ValueError("Expected a single video, got a list of video tensors/arrays.") - return _normalize_frames(video) + return np.stack(_normalize_frames(video), axis=0) raise ValueError("Unsupported list contents for video payload.") raise ValueError(f"Unsupported video payload type: {type(video)}") -def _coerce_audio_to_waveform(audio: Any) -> torch.Tensor: - """Convert an audio payload into a 2-channel CPU float tensor for LTX2 export.""" +def _coerce_audio_to_numpy(audio: Any) -> np.ndarray: + """Convert an audio payload into a float32 numpy array for muxing.""" if isinstance(audio, torch.Tensor): - waveform = audio.detach().cpu() + arr = audio.detach().cpu().float().numpy() elif isinstance(audio, np.ndarray): - waveform = torch.from_numpy(audio) + arr = audio elif isinstance(audio, list): - waveform = torch.tensor(audio) + arr = np.array(audio) else: raise ValueError(f"Unsupported audio payload type: {type(audio)}") - waveform = waveform.squeeze() - - if waveform.ndim == 0: + arr = np.squeeze(arr) + if arr.ndim == 0: raise ValueError("Audio payload must contain at least one sample.") - if waveform.ndim == 1: - waveform = waveform.unsqueeze(0) - elif waveform.ndim == 2: - if waveform.shape[0] in (1, 2): - pass - elif waveform.shape[1] in (1, 2): - waveform = waveform.transpose(0, 1) - else: - raise ValueError(f"Unsupported audio payload shape: {tuple(waveform.shape)}") - else: - raise ValueError(f"Unsupported audio payload rank: {waveform.ndim}") - - if waveform.shape[0] == 1: - waveform = waveform.repeat(2, 1) - elif waveform.shape[0] != 2: - raise ValueError(f"Expected mono or stereo audio, got shape {tuple(waveform.shape)}") - - return waveform.float().contiguous() + return arr.astype(np.float32) -def _encode_video_bytes(video: Any, fps: int, audio: Any | None = None, audio_sample_rate: int | None = None) -> bytes: +def _encode_video_bytes( + video: Any, + fps: int, + audio: Any | None = None, + audio_sample_rate: int | None = None, + video_codec_options: dict[str, str] | None = None, +) -> bytes: """Encode a video payload into MP4 bytes, optionally muxing audio.""" - try: - from diffusers.utils import export_to_video - except ImportError as exc: # pragma: no cover - optional dependency - raise ImportError("diffusers is required for export_to_video.") from exc - - frames = _coerce_video_to_frames(video) - if not frames: + frames_np = _coerce_video_to_array(video) + if frames_np.size == 0: raise ValueError("No frames found to encode.") + if frames_np.ndim == 4 and frames_np.shape[-1] == 4: + frames_np = frames_np[..., :3] - tmp_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) - tmp_file.close() - try: - if audio is not None: - from diffusers.pipelines.ltx2.export_utils import encode_video as encode_ltx2_video + if audio is None: + from diffusers.utils import export_to_video - frames_np = np.stack(frames, axis=0) - if frames_np.ndim == 4 and frames_np.shape[-1] == 4: - frames_np = frames_np[..., :3] - frames_np = np.clip(frames_np, 0.0, 1.0) - frames_u8 = (frames_np * 255).round().clip(0, 255).astype("uint8") - video_tensor = torch.from_numpy(frames_u8) - encode_ltx2_video( - video_tensor, - fps=fps, - audio=_coerce_audio_to_waveform(audio), - audio_sample_rate=audio_sample_rate, - output_path=tmp_file.name, - ) - else: - export_to_video(frames, tmp_file.name, fps=fps) - with open(tmp_file.name, "rb") as f: - return f.read() - finally: + tmp_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) + tmp_file.close() try: - os.remove(tmp_file.name) - except OSError: - pass + export_to_video(list(frames_np), tmp_file.name, fps=fps) + with open(tmp_file.name, "rb") as f: + return f.read() + finally: + try: + os.remove(tmp_file.name) + except OSError: + pass + + from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes + + if frames_np.dtype == np.uint8: + frames_u8 = frames_np + else: + frames_np = np.clip(frames_np, 0.0, 1.0) + frames_np *= 255.0 + frames_u8 = np.round(frames_np).astype(np.uint8) + + # Ensure contiguous memory layout for faster PyAV muxing + frames_u8 = np.ascontiguousarray(frames_u8) + + audio_np = _coerce_audio_to_numpy(audio) if audio is not None else None + + return mux_video_audio_bytes( + frames_u8, + audio_np, + fps=float(fps), + audio_sample_rate=audio_sample_rate or 24000, + video_codec_options=video_codec_options, + ) -def encode_video_base64(video: Any, fps: int, audio: Any | None = None, audio_sample_rate: int | None = None) -> str: +def encode_video_base64( + video: Any, + fps: int, + audio: Any | None = None, + audio_sample_rate: int | None = None, + video_codec_options: dict[str, str] | None = None, +) -> str: """Encode a video (frames/array/tensor) to base64 MP4.""" - video_bytes = _encode_video_bytes(video, fps=fps, audio=audio, audio_sample_rate=audio_sample_rate) + video_bytes = _encode_video_bytes( + video, fps=fps, audio=audio, audio_sample_rate=audio_sample_rate, video_codec_options=video_codec_options + ) return base64.b64encode(video_bytes).decode("utf-8") diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index 5768c3b6d..2a63333c8 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -227,6 +227,10 @@ class OmniDiffusionSamplingParams: frame_rate: float | None = None # Floating-point rate used by the diffusion model when it differs from `fps`. height_not_provided: bool = False width_not_provided: bool = False + enable_frame_interpolation: bool = False + frame_interpolation_exp: int = 1 + frame_interpolation_scale: float = 1.0 + frame_interpolation_model_path: str | None = None # Timesteps timesteps: torch.Tensor | None = None diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py index ca3ba271a..01acf5bb0 100644 --- a/vllm_omni/outputs.py +++ b/vllm_omni/outputs.py @@ -58,6 +58,10 @@ class OmniRequestOutput: images: list[Image.Image] = field(default_factory=list) prompt: OmniPromptType | None = None latents: torch.Tensor | None = None + trajectory_timesteps: list[torch.Tensor] | None = None + trajectory_latents: torch.Tensor | None = None + trajectory_log_probs: torch.Tensor | None = None + trajectory_decoded: list[torch.Tensor] | list[Image.Image] | None = None metrics: dict[str, Any] = field(default_factory=dict) _multimodal_output: dict[str, Any] = field(default_factory=dict) _custom_output: dict[str, Any] = field(default_factory=dict) @@ -68,6 +72,9 @@ class OmniRequestOutput: # memory usage info peak_memory_mb: float = 0.0 + # error handling + error: str | None = None + @classmethod def from_pipeline( cls, @@ -101,6 +108,10 @@ class OmniRequestOutput: prompt: OmniPromptType | None = None, metrics: dict[str, Any] | None = None, latents: torch.Tensor | None = None, + trajectory_timesteps: list[torch.Tensor] | None = None, + trajectory_latents: torch.Tensor | None = None, + trajectory_log_probs: torch.Tensor | None = None, + trajectory_decoded: list[torch.Tensor] | list[Image.Image] | None = None, multimodal_output: dict[str, Any] | None = None, custom_output: dict[str, Any] | None = None, final_output_type: str = "image", @@ -115,6 +126,10 @@ class OmniRequestOutput: prompt: The prompt used metrics: Generation metrics latents: Optional latent tensors + trajectory_timesteps: Optional sampled timesteps for trajectory outputs + trajectory_latents: Optional latent trajectory tensor + trajectory_log_probs: Optional per-step log probabilities + trajectory_decoded: Optional decoded trajectory frames multimodal_output: Optional multimodal output dict custom_output: Optional custom output dict (e.g. latent trajectories, prompt embeds) stage_durations: Optional stage durations (execution time of each stage) dict @@ -129,6 +144,10 @@ class OmniRequestOutput: images=images, prompt=prompt, latents=latents, + trajectory_timesteps=trajectory_timesteps, + trajectory_latents=trajectory_latents, + trajectory_log_probs=trajectory_log_probs, + trajectory_decoded=trajectory_decoded, metrics=metrics or {}, _multimodal_output=multimodal_output or {}, _custom_output=custom_output or {},