[cherry-pick releasev0.18.0post1] Codex/release v0.18.0.post1 cherry-pick #2735/#2736/#2555 (#2813)

Signed-off-by: samithuang <285365963@qq.com>
Signed-off-by: David Chen <530634352@qq.com>
Signed-off-by: Samit <285365963@qq.com>
Signed-off-by: david6666666 <530634352@qq.com>
Co-authored-by: Samit <285365963@qq.com>

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

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`.

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
 

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)
 
 

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()

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.

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)

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},
     }
+
+

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
 

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,

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

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
 
 

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
 

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
 

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
 

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"]

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: <model_path>/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)

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,

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()

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)",

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

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:

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),
     }

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,

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

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")

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

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 {},