[feat] Streaming diffusion video generation output (#3737)

Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
Signed-off-by: Zeyu Huang | 黃澤宇 <11222265+fhfuih@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

tests/entrypoints/openai_api/test_video_api_utils.py

@@ -48,6 +48,107 @@ def test_encode_video_bytes_exports_frames_without_interpolation(monkeypatch):
     assert mux_calls[0]["audio"] is None
 
 
+def test_fragmented_mp4_video_encoder_reuses_single_muxer(monkeypatch):
+    muxers = []
+
+    class FakeFragmentedMP4Muxer:
+        def __init__(self, *, width, height, fps, video_codec_options=None):
+            self.calls = []
+            muxers.append(
+                {
+                    "width": width,
+                    "height": height,
+                    "fps": fps,
+                    "video_codec_options": video_codec_options,
+                    "instance": self,
+                }
+            )
+
+        def mux_video_frames(self, frames):
+            self.calls.append(frames)
+            return f"fragment-{len(self.calls)}".encode()
+
+        def close(self):
+            return b"trailer"
+
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.utils.media_utils.FragmentedMP4Muxer",
+        FakeFragmentedMP4Muxer,
+    )
+
+    encoder = video_api_utils.FragmentedMP4VideoEncoder(
+        fps=12,
+        video_codec_options={"preset": "ultrafast"},
+    )
+    assert encoder.encode([np.zeros((4, 6, 3), dtype=np.float32)]) == b"fragment-1"
+    assert encoder.encode([np.ones((4, 6, 3), dtype=np.float32)]) == b"fragment-2"
+    assert encoder.close() == b"trailer"
+
+    assert len(muxers) == 1
+    assert muxers[0]["width"] == 6
+    assert muxers[0]["height"] == 4
+    assert muxers[0]["fps"] == 12.0
+    assert muxers[0]["video_codec_options"] == {"preset": "ultrafast"}
+    assert len(muxers[0]["instance"].calls) == 2
+
+
+def test_create_streaming_video_encoder_selects_requested_format():
+    assert isinstance(
+        video_api_utils.create_streaming_video_encoder(output_format="m4s", fps=12),
+        video_api_utils.FragmentedMP4VideoEncoder,
+    )
+
+
+def test_finalize_streaming_mp4_bytes_produces_progressive_mp4():
+    """Fragment MP4 chunks are remuxed into decodable progressive MP4 bytes."""
+    import numpy as np
+
+    from vllm_omni.diffusion.utils.media_utils import FragmentedMP4Muxer, finalize_streaming_video_bytes
+
+    def _read_mp4_video_info(mp4_bytes: bytes) -> tuple[int, float, float | None, int, int]:
+        import io
+
+        import av
+
+        with av.open(io.BytesIO(mp4_bytes), mode="r", format="mp4") as container:
+            stream = container.streams.video[0]
+            frame_count = sum(1 for _ in container.decode(stream))
+            assert stream.average_rate is not None
+            fps = float(stream.average_rate)
+            duration_sec = None
+            if stream.duration is not None:
+                assert stream.time_base is not None
+                duration_sec = float(stream.duration * stream.time_base)
+            return frame_count, fps, duration_sec, stream.width, stream.height
+
+    width = 32
+    height = 32
+    fps = 16.0
+    input_frame_count = 2
+
+    muxer = FragmentedMP4Muxer(width=width, height=height, fps=fps)
+    frames = np.zeros((input_frame_count, height, width, 3), dtype=np.uint8)
+    streamed = muxer.mux_video_frames(frames) + muxer.close()
+
+    finalized = finalize_streaming_video_bytes(streamed, input_format="m4s", fps=fps)
+    assert finalized
+    assert finalized != streamed
+
+    streamed_info = _read_mp4_video_info(streamed)
+    final_info = _read_mp4_video_info(finalized)
+    expected_duration = input_frame_count / fps
+
+    assert streamed_info[0] == input_frame_count
+    assert final_info[0] == input_frame_count
+    assert streamed_info[1] == pytest.approx(fps)
+    assert final_info[1] == pytest.approx(fps)
+    assert streamed_info[3:] == (width, height)
+    assert final_info[3:] == (width, height)
+
+    assert final_info[2] == pytest.approx(expected_duration, rel=0.05)
+    assert streamed_info[2] == pytest.approx(expected_duration, rel=0.05)
+
+
 def test_rife_model_inference_runs_on_dummy_tensors():
     model = rife_interpolator.Model().eval()
     img0 = torch.rand(1, 3, 32, 32)