[5/N] Schedule chunk transfers as their own stage (#7406)

Signed-off-by: specture724 <specture724@gmail.com>

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -2859,8 +2859,10 @@ def test_decode_to_mp4_batches_consumer_transfers(monkeypatch):
             self.kwargs = kwargs
             self.__class__.instances.append(self)
 
-        def push(self, frames):
+        def push(self, frames, *, on_consumed=None):
             self.pushes.append(np.array(frames, copy=True))
+            if on_consumed is not None:
+                on_consumed()
 
         def finish(self):
             return b"mp4"
@@ -2986,8 +2988,10 @@ def test_request_video_codec_options_reach_the_preencoded_mp4_encoder(monkeypatc
             self.kwargs = kwargs
             self.__class__.instances.append(self)
 
-        def push(self, frames):
+        def push(self, frames, *, on_consumed=None):
             del frames
+            if on_consumed is not None:
+                on_consumed()
 
         def finish(self):
             return b"mp4"

tests/diffusion/models/wan2_2/test_wan22_chunked_mp4.py

@@ -138,13 +138,13 @@ def test_request_batch_threshold_controls_transfers_without_losing_frames(monkey
     from vllm_omni.diffusion.utils import chunked_video
 
     transfers = []
-    quantize = chunked_video.chunk_to_uint8_frames
+    quantize = chunked_video.quantize_chunk
 
     def capture(chunk, value_range):
         transfers.append(chunk.shape[2])
         return quantize(chunk, value_range)
 
-    monkeypatch.setattr(chunked_video, "chunk_to_uint8_frames", capture)
+    monkeypatch.setattr(chunked_video, "quantize_chunk", capture)
     params = SimpleNamespace(extra_args={"preencode_mp4": True, "preencode_batch_frames": value})
     vae = _FakeWanVAE(batch=1, chunks=4, frames_per_chunk=3, height=16, width=32)
     videos = decode_to_mp4(vae, torch.zeros(1), fps=24, batch_frames=resolve_wan_preencode_batch_frames(params))

tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from types import SimpleNamespace
+from typing import Any
 from unittest.mock import MagicMock, patch
 
 import numpy as np
@@ -43,7 +44,7 @@ def test_s2v_constructor_preserves_unshifted_endpoints(configured_shift: float |
 
 
 def _make_s2v_sampling(**overrides):
-    values: dict[str, object] = {
+    values: dict[str, Any] = {
         "height": 16,
         "width": 16,
         "num_frames": 8,
@@ -749,13 +750,13 @@ def test_s2v_preencode_returns_playable_mp4_bytes_per_request(monkeypatch, batch
     from vllm_omni.diffusion.utils import chunked_video
 
     transfers = []
-    quantize = chunked_video.chunk_to_uint8_frames
+    quantize = chunked_video.quantize_chunk
 
     def capture(chunk, value_range):
         transfers.append(chunk.shape[2])
         return quantize(chunk, value_range)
 
-    monkeypatch.setattr(chunked_video, "chunk_to_uint8_frames", capture)
+    monkeypatch.setattr(chunked_video, "quantize_chunk", capture)
 
     with patch("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.current_omni_platform") as platform:
         platform.is_available.return_value = False