[4/N] Unify worker-side MP4 encoding and wire Wan into it (#7048)

Signed-off-by: specture724 <specture724@gmail.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
authored and GitHub committed 2026-09-17 23:33:56 +08:00
1 parent 882bf8bc95
commit 2ab5d1701a
12 files changed
+1177 -232

No files matched your search

@@ -2852,7 +2852,7 @@ def test_decode_to_mp4_batches_consumer_transfers(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
class FakeEncoder:
instances = []
instances: list["FakeEncoder"] = []
def __init__(self, **kwargs):
self.pushes = []
@@ -2873,6 +2873,8 @@ def test_decode_to_mp4_batches_consumer_transfers(monkeypatch):
return torch.zeros(1, 1, 2)
class FakeVideoVAE:
chunk_value_range = (0.0, 1.0)
def decode_with_chunks(self, latent, *, on_chunk):
for value in (0.0, 0.25, 0.5):
on_chunk(torch.full((1, 3, 3, 2, 2), value))
@@ -2884,7 +2886,7 @@ def test_decode_to_mp4_batches_consumer_transfers(monkeypatch):
pipeline.device = torch.device("cpu")
monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_uses_manual_component_offload", lambda self, component: False)
monkeypatch.setattr(
"vllm_omni.diffusion.utils.media_utils.ChunkedMP4Encoder",
"vllm_omni.diffusion.utils.chunked_video.ChunkedMP4Encoder",
FakeEncoder,
)
@@ -2904,12 +2906,81 @@ def test_decode_to_mp4_batches_consumer_transfers(monkeypatch):
assert encoder.pushes[1].shape == (3, 2, 2, 3)
def _peer_rank_preencode_pipeline(monkeypatch):
"""An H3 pipeline whose video VAE is a distributed rank that owns no output.
Every rank of the VAE group must call ``decode_with_chunks`` to keep the
temporal collectives in lockstep, but only the owner receives chunks, so a
peer's callback is never invoked.
"""
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
class FakeAudioVAE:
def decode_latent(self, latent):
return torch.zeros(1, 1, 2)
class PeerRankVideoVAE:
chunk_value_range = (0.0, 1.0)
def decode_with_chunks(self, latent, *, on_chunk):
del latent, on_chunk
pipeline = object.__new__(mod.MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.audio_vae = FakeAudioVAE()
pipeline.video_vae = PeerRankVideoVAE()
pipeline.device = torch.device("cpu")
pipeline.od_config = SimpleNamespace()
monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_uses_manual_component_offload", lambda self, component: False)
return pipeline
def test_peer_vae_rank_returns_no_preencoded_output_instead_of_failing(monkeypatch):
"""A rank that owns no output must not turn its silence into a request error."""
pipeline = _peer_rank_preencode_pipeline(monkeypatch)
output = pipeline.decode_to_mp4(torch.zeros(1), torch.zeros(1), height=2, width=2)
assert output == b""
def test_peer_vae_rank_post_decode_reaches_post_processing(monkeypatch):
"""The real H3 post_decode chain, not just the shared consumer, tolerates a peer rank."""
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
from vllm_omni.diffusion.worker.utils import StepRequestState
pipeline = _peer_rank_preencode_pipeline(monkeypatch)
monkeypatch.setattr(pipeline, "_unpack_denoised_rows", lambda *args, **kwargs: (torch.zeros(1), torch.zeros(1)))
state = StepRequestState(
request_id="peer-rank",
sampling=SimpleNamespace(num_outputs_per_prompt=1),
prompt="a prompt",
)
state.latents = torch.zeros(1)
state.extra[mod._STEP_BRANCH] = object()
state.extra[mod._STEP_AUDIO_ROWS] = torch.zeros(1)
state.extra[mod._STEP_SHAPE] = {
"latent_t": 1,
"latent_h": 1,
"latent_w": 1,
"audio_t": 1,
"height": 2,
"width": 2,
"preencode_mp4": True,
}
output = pipeline.post_decode(state)
assert output.output == (b"", None)
assert mod._minimax_h3_post_process(output.output)["video"] == [b""]
def test_request_video_codec_options_reach_the_preencoded_mp4_encoder(monkeypatch):
"""A client's encoder options must survive the worker-side pre-encode path."""
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
class FakeEncoder:
instances = []
instances: list["FakeEncoder"] = []
def __init__(self, **kwargs):
self.kwargs = kwargs
@@ -2929,6 +3000,8 @@ def test_request_video_codec_options_reach_the_preencoded_mp4_encoder(monkeypatc
return torch.zeros(1, 1, 2)
class FakeVideoVAE:
chunk_value_range = (0.0, 1.0)
def decode_with_chunks(self, latent, *, on_chunk):
on_chunk(torch.zeros(1, 3, 1, 2, 2))
@@ -2938,7 +3011,7 @@ def test_request_video_codec_options_reach_the_preencoded_mp4_encoder(monkeypatc
pipeline.video_vae = FakeVideoVAE()
pipeline.device = torch.device("cpu")
monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_uses_manual_component_offload", lambda self, component: False)
monkeypatch.setattr("vllm_omni.diffusion.utils.media_utils.ChunkedMP4Encoder", FakeEncoder)
monkeypatch.setattr("vllm_omni.diffusion.utils.chunked_video.ChunkedMP4Encoder", FakeEncoder)
pipeline.decode_to_mp4(
torch.zeros(1),
@@ -0,0 +1,165 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import io
from types import SimpleNamespace
import av
import pytest
import torch
from vllm_omni.diffusion.models.wan2_2.chunked_mp4 import (
WAN_DEFAULT_OUTPUT_FPS,
resolve_wan_output_fps,
resolve_wan_preencode_batch_frames,
resolve_wan_preencode_mp4,
resolve_wan_video_codec_options,
wan_preencoded_mp4_payload,
)
from vllm_omni.diffusion.utils.chunked_video import decode_to_mp4
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
class _FakeWanVAE:
"""Publish temporal chunks the way the Wan VAE seam does."""
chunk_value_range = (-1.0, 1.0)
def __init__(self, *, batch: int, chunks: int, frames_per_chunk: int, height: int, width: int):
self.batch = batch
self.chunks = chunks
self.frames_per_chunk = frames_per_chunk
self.height = height
self.width = width
self.chunk_calls = 0
def decode_with_chunks(self, latents, *, on_chunk):
del latents
assert on_chunk is not None
for index in range(self.chunks):
self.chunk_calls += 1
value = -1.0 + 2.0 * index / max(self.chunks - 1, 1)
on_chunk(torch.full((self.batch, 3, self.frames_per_chunk, self.height, self.width), value))
def _decoded_frame_count(data: bytes) -> int:
with av.open(io.BytesIO(data)) as container:
return sum(1 for _ in container.decode(video=0))
def test_decode_wan_latents_to_mp4_emits_one_playable_mp4_per_batch_entry():
vae = _FakeWanVAE(batch=2, chunks=6, frames_per_chunk=2, height=16, width=32)
videos = decode_to_mp4(vae, torch.zeros(1), fps=24, batch_frames=4)
assert len(videos) == 2
for data in videos:
assert data[:4] and _decoded_frame_count(data) == 12
def test_decode_wan_latents_to_mp4_coalesces_transfers_without_losing_frames():
"""batch_frames only changes how many transfers carry the same frames."""
unbatched = _FakeWanVAE(batch=1, chunks=8, frames_per_chunk=1, height=16, width=32)
batched = _FakeWanVAE(batch=1, chunks=8, frames_per_chunk=1, height=16, width=32)
per_chunk = decode_to_mp4(unbatched, torch.zeros(1), fps=24, batch_frames=1)
coalesced = decode_to_mp4(batched, torch.zeros(1), fps=24, batch_frames=4)
assert _decoded_frame_count(per_chunk[0]) == _decoded_frame_count(coalesced[0]) == 8
def test_decode_wan_latents_to_mp4_returns_nothing_for_a_rank_without_output():
class SilentVAE:
chunk_value_range = (-1.0, 1.0)
def decode_with_chunks(self, latents, *, on_chunk):
del latents, on_chunk
assert decode_to_mp4(SilentVAE(), torch.zeros(1), fps=24) == []
def test_decode_wan_latents_to_mp4_rejects_a_non_positive_batch():
with pytest.raises(ValueError, match="batch_frames"):
decode_to_mp4(
_FakeWanVAE(batch=1, chunks=1, frames_per_chunk=1, height=16, width=16), None, fps=24, batch_frames=0
)
def test_decode_wan_latents_to_mp4_rejects_a_vae_without_the_capability():
class PlainVAE:
chunk_value_range = (-1.0, 1.0)
def decode(self, z, return_dict=True):
del z, return_dict
with pytest.raises(TypeError, match="chunked VAE decode capability"):
decode_to_mp4(PlainVAE(), torch.zeros(1), fps=24)
def test_preencode_flag_is_off_unless_the_request_asks_for_it():
assert resolve_wan_preencode_mp4(SimpleNamespace(extra_args=None), output_type="np") is False
assert resolve_wan_preencode_mp4(SimpleNamespace(extra_args={}), output_type="np") is False
assert resolve_wan_preencode_mp4(SimpleNamespace(extra_args={"preencode_mp4": True}), output_type="np") is True
def test_preencode_rejects_combinations_it_cannot_serve():
interpolating = SimpleNamespace(extra_args={"preencode_mp4": True}, enable_frame_interpolation=True)
with pytest.raises(ValueError, match="enable_frame_interpolation"):
resolve_wan_preencode_mp4(interpolating, output_type="np")
with pytest.raises(ValueError, match="output_type"):
resolve_wan_preencode_mp4(SimpleNamespace(extra_args={"preencode_mp4": True}), output_type="pil")
def test_output_fps_mirrors_the_serving_fallback():
assert resolve_wan_output_fps(SimpleNamespace(fps=16)) == 16
assert resolve_wan_output_fps(SimpleNamespace(fps=[16])) == 16
assert resolve_wan_output_fps(SimpleNamespace(fps=None)) == WAN_DEFAULT_OUTPUT_FPS
assert resolve_wan_output_fps(SimpleNamespace()) == WAN_DEFAULT_OUTPUT_FPS
def test_codec_options_reach_the_encoder_as_strings():
params = SimpleNamespace(extra_args={"video_codec_options": {"preset": "ultrafast", "threads": 0}})
assert resolve_wan_video_codec_options(params) == {"preset": "ultrafast", "threads": "0"}
assert resolve_wan_video_codec_options(SimpleNamespace(extra_args={})) == {"preset": "ultrafast", "threads": "0"}
assert resolve_wan_video_codec_options(SimpleNamespace(extra_args={"video_codec_options": None})) is None
assert resolve_wan_video_codec_options(SimpleNamespace(extra_args={"video_codec_options": {}})) == {}
def test_preencoded_payload_passes_bytes_through_and_ignores_tensors():
assert wan_preencoded_mp4_payload(torch.zeros(1, 3, 2, 4, 4)) is None
assert wan_preencoded_mp4_payload([b"one", b"two"]) == {"payload": {"video": [b"one", b"two"]}, "metadata": {}}
assert wan_preencoded_mp4_payload(b"one") == {"payload": {"video": [b"one"]}, "metadata": {}}
@pytest.mark.parametrize("value", [1, 5, 33])
def test_request_batch_threshold_controls_transfers_without_losing_frames(monkeypatch, value):
from vllm_omni.diffusion.utils import chunked_video
transfers = []
quantize = chunked_video.chunk_to_uint8_frames
def capture(chunk, value_range):
transfers.append(chunk.shape[2])
return quantize(chunk, value_range)
monkeypatch.setattr(chunked_video, "chunk_to_uint8_frames", 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))
assert transfers == {1: [3, 3, 3, 3], 5: [6, 6], 33: [12]}[value]
assert _decoded_frame_count(videos[0]) == 12
@pytest.mark.parametrize("value", [0, -1, True, 1.5, "17", None])
def test_request_batch_threshold_rejects_invalid_values(value):
with pytest.raises(ValueError, match="preencode_batch_frames"):
resolve_wan_preencode_batch_frames(SimpleNamespace(extra_args={"preencode_batch_frames": value}))
def test_request_batch_threshold_preserves_model_defaults():
params = SimpleNamespace(extra_args=None)
assert resolve_wan_preencode_batch_frames(params) == 17
assert resolve_wan_preencode_batch_frames(params, default=1) == 1
@@ -647,3 +647,245 @@ def test_s2v_pipeline_hsdp_forward_complete_process():
assert video.shape[1] == 3 # channels
assert audio_waveform is not None
assert audio_sr == 16000
def _make_s2v_preencode_pipeline() -> Wan22S2VPipeline:
"""Build the same stub pipeline the batching test drives, for preencode runs."""
pipeline = object.__new__(Wan22S2VPipeline)
nn.Module.__init__(pipeline)
pipeline.device = torch.device("cpu")
pipeline.transformer = MagicMock()
pipeline.transformer.dtype = torch.float32
# A fresh iterator per call: the clip loop reads this once per clip.
pipeline.transformer.parameters.side_effect = lambda: iter([torch.zeros(1, dtype=torch.float32)])
pipeline.transformer.casual_audio_encoder = None
pipeline.transformer.encode_audio.side_effect = lambda audio, _motion: {"audio_emb": audio}
pipeline.vae = MagicMock()
pipeline.vae.dtype = torch.float32
# The shared consumer reads the published pixel range off the VAE.
pipeline.vae.chunk_value_range = (-1.0, 1.0)
# One decoded clip: [B, C, T, H, W] for a batch of four (2 requests x 2 outputs).
pipeline.vae.decode.return_value = (torch.zeros(4, 3, 8, 16, 16),)
pipeline.od_config = SimpleNamespace(
enable_cpu_offload=False,
parallel_config=SimpleNamespace(use_hsdp=False),
)
pipeline.scheduler = MagicMock(timesteps=torch.tensor([1.0]))
pipeline._flow_shift = 3.0
pipeline.vae_scale_factor_spatial = 8
pipeline.resolution_divisor = 16
pipeline.motion_frames = 7
pipeline.drop_first_motion = True
pipeline._DEFAULT_INFER_FRAMES = 8
pipeline._guidance_scale = None
pipeline._num_timesteps = None
pipeline.check_inputs = lambda *args, **kwargs: None
pipeline.encode_prompt = MagicMock(return_value=(torch.zeros(4, 2, 3), torch.ones(4, 2, 3)))
pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 1, 2, 8), 1, 8))
pipeline.encode_ref_image = MagicMock(return_value=torch.zeros(1, 16, 1, 2, 2))
pipeline.prepare_motion_latents = MagicMock(
side_effect=lambda pixels, **_: torch.zeros(pixels.shape[0], 16, 2, 2, 2)
)
pipeline._denormalize_latents = lambda latents: latents
pipeline.diffuse = MagicMock(side_effect=lambda **kwargs: kwargs["latents"])
return pipeline
def _make_s2v_preencode_batch(audio_a, audio_b, **sampling_overrides) -> DiffusionRequestBatch:
image = PIL.Image.new("RGB", (16, 16))
return DiffusionRequestBatch(
requests=[
SimpleNamespace(
request_id="a",
prompt={"prompt": "first", "multi_modal_data": {"image": image, "audio": audio_a}},
sampling_params=_make_s2v_sampling(
num_outputs_per_prompt=2,
latents=torch.zeros(2, 16, 2, 2, 2),
output_type="np",
**sampling_overrides,
),
),
SimpleNamespace(
request_id="b",
prompt={"prompt": "second", "multi_modal_data": {"image": image, "audio": audio_b}},
sampling_params=_make_s2v_sampling(
num_outputs_per_prompt=2,
latents=torch.ones(2, 16, 2, 2, 2),
output_type="np",
**sampling_overrides,
),
),
]
)
def _decode_mp4(data: bytes):
"""Return (video frame count, first audio samples) for one encoded container."""
import io
import av
with av.open(io.BytesIO(data)) as container:
frames = sum(1 for _ in container.decode(video=0))
with av.open(io.BytesIO(data)) as container:
assert container.streams.audio, "the worker must mux the waveform into the container"
samples = np.concatenate(
[frame.to_ndarray().reshape(-1) for frame in container.decode(audio=0)][:4],
)
return frames, samples
@pytest.mark.parametrize("batch_frames", [1, 1000])
def test_s2v_preencode_returns_playable_mp4_bytes_per_request(monkeypatch, batch_frames) -> None:
"""The clip loop hands finished clips to the encoder instead of concatenating."""
pipeline = _make_s2v_preencode_pipeline()
# Two clips, so the autoregressive motion feedback runs between pushes.
pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 1, 2, 8), 2, 16))
audio_a = np.zeros(16000, dtype=np.float32)
audio_b = np.full(16000, 0.5, dtype=np.float32)
batch = _make_s2v_preencode_batch(
audio_a, audio_b, extra_args={"preencode_mp4": True, "preencode_batch_frames": batch_frames}
)
from vllm_omni.diffusion.utils import chunked_video
transfers = []
quantize = chunked_video.chunk_to_uint8_frames
def capture(chunk, value_range):
transfers.append(chunk.shape[2])
return quantize(chunk, value_range)
monkeypatch.setattr(chunked_video, "chunk_to_uint8_frames", capture)
with patch("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.current_omni_platform") as platform:
platform.is_available.return_value = False
outputs = pipeline.forward(batch)
assert len(outputs) == 2
assert pipeline.vae.decode.call_count == 2, "both clips must reach the encoder"
assert len(transfers) == (2 if batch_frames == 1 else 1)
decoded = []
for request_output in outputs:
# Two outputs per prompt, each already a complete container.
assert isinstance(request_output.output, list)
assert len(request_output.output) == 2
for data in request_output.output:
assert isinstance(data, bytes)
frames, samples = _decode_mp4(data)
assert frames > 0
decoded.append(samples)
# Four entries: request a's waveform for its two outputs, then request b's.
# A silent and a non-silent waveform must not land on the same request.
assert np.abs(decoded[0]).max() == pytest.approx(0.0, abs=1e-3)
assert np.abs(decoded[1]).max() == pytest.approx(0.0, abs=1e-3)
assert np.abs(decoded[2]).max() > 0.1
assert np.abs(decoded[3]).max() > 0.1
def test_s2v_preencode_aborts_encoders_when_a_later_clip_fails(monkeypatch) -> None:
"""A failure after the first push must not strand encoder worker threads."""
pipeline = _make_s2v_preencode_pipeline()
pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 1, 2, 8), 2, 16))
batch = _make_s2v_preencode_batch(
np.zeros(16000, dtype=np.float32),
np.ones(16000, dtype=np.float32),
extra_args={"preencode_mp4": True, "preencode_batch_frames": 1},
)
sessions = []
from vllm_omni.diffusion.utils.chunked_video import ChunkedVideoMP4Session
def capture_session(**kwargs):
session = ChunkedVideoMP4Session(**kwargs)
sessions.append(session)
return session
def fail_on_second_clip(**kwargs):
if pipeline.vae.decode.call_count:
raise RuntimeError("second clip failed")
return kwargs["latents"]
pipeline.diffuse = fail_on_second_clip
monkeypatch.setattr("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.ChunkedVideoMP4Session", capture_session)
with patch("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.current_omni_platform") as platform:
platform.is_available.return_value = False
with pytest.raises(RuntimeError, match="second clip failed"):
pipeline.forward(batch)
assert len(sessions) == 1
assert len(sessions[0]._encoders) == 4
assert all(not encoder._thread.is_alive() for encoder in sessions[0]._encoders)
def test_s2v_preencode_skips_encoding_on_a_vae_patch_parallel_peer(monkeypatch) -> None:
"""A peer still receives every clip for the motion loop but encodes none of them."""
pipeline = _make_s2v_preencode_pipeline()
pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 1, 2, 8), 2, 16))
# A peer's patch-parallel decode returns an empty placeholder.
pipeline.vae.decode.return_value = (torch.empty(0),)
pipeline.vae._vae_pp_group = object()
broadcasts = []
monkeypatch.setattr(
"torch.distributed.broadcast",
lambda tensor, src, group: broadcasts.append(tensor.zero_()),
)
batch = _make_s2v_preencode_batch(
np.zeros(16000, dtype=np.float32),
np.ones(16000, dtype=np.float32),
extra_args={"preencode_mp4": True, "preencode_batch_frames": 1},
)
sessions = []
from vllm_omni.diffusion.utils.chunked_video import ChunkedVideoMP4Session
def capture_session(**kwargs):
session = ChunkedVideoMP4Session(**kwargs)
sessions.append(session)
return session
monkeypatch.setattr("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.ChunkedVideoMP4Session", capture_session)
with patch("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.current_omni_platform") as platform:
platform.is_available.return_value = False
outputs = pipeline.forward(batch)
assert len(broadcasts) == 2, "the peer must still join the broadcast for both clips"
# Once for the initial motion latents, once more from the peer's broadcast clip.
assert pipeline.prepare_motion_latents.call_count == 2, "the motion loop still consumes the peer's frames"
assert sessions[0]._encoders == []
assert [request_output.output for request_output in outputs] == [[], []]
@pytest.mark.parametrize("output_type", ["pil", "pt", "latent"])
def test_s2v_preencode_rejects_request_output_types_it_cannot_serve(output_type) -> None:
"""Pre-encoding returns MP4 bytes, so a request asking for frames must fail."""
pipeline = _make_s2v_preencode_pipeline()
batch = _make_s2v_preencode_batch(
np.zeros(16000, dtype=np.float32),
np.ones(16000, dtype=np.float32),
extra_args={"preencode_mp4": True},
)
for request in batch.requests:
request.sampling_params.output_type = output_type
with pytest.raises(ValueError, match="output_type"):
pipeline.forward(batch)
def test_s2v_preencode_keeps_the_full_decode_path_untouched() -> None:
"""Without the flag the loop still returns the (video, audio, rate) tuple."""
pipeline = _make_s2v_preencode_pipeline()
audio_a = np.zeros(16000, dtype=np.float32)
audio_b = np.ones(16000, dtype=np.float32)
batch = _make_s2v_preencode_batch(audio_a, audio_b)
with patch("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.current_omni_platform") as platform:
platform.is_available.return_value = False
outputs = pipeline.forward(batch)
assert outputs[0].output[0].shape[0] == 2
np.testing.assert_array_equal(outputs[0].output[1], audio_a)
np.testing.assert_array_equal(outputs[1].output[1], audio_b)