diff --git a/docs/serving/videos_api.md b/docs/serving/videos_api.md index 1152ebce3..f48f2ca99 100644 --- a/docs/serving/videos_api.md +++ b/docs/serving/videos_api.md @@ -251,6 +251,7 @@ These `extra_params` control how the server turns decoded frames into MP4 bytes. | Field | Type | Default | Description | | --- | --- | --- | --- | | `preencode_mp4` | boolean | false | Encode the MP4 on the worker while the VAE is still decoding, instead of after the full video is materialized | +| `preencode_batch_frames` | positive integer | 17 (H3, Wan T2V/I2V); 1 (Wan S2V) | Minimum accumulated frames per worker transfer/encoding batch; used only with `preencode_mp4=true` | | `video_codec_options` | object | null | Encoder options passed through to the H.264 encoder, such as `{"preset": "ultrafast", "threads": "0"}` | With `preencode_mp4` enabled, each committed VAE chunk leaves the accelerator and @@ -261,16 +262,29 @@ unchanged: the same complete MP4, byte-for-byte equivalent frames. ```bash curl -X POST http://localhost:8091/v1/videos/sync \ -F "prompt=A small robot walking through a neon city" \ - -F 'extra_params={"preencode_mp4": true, "video_codec_options": {"preset": "ultrafast"}}' \ + -F 'extra_params={"preencode_mp4": true, "preencode_batch_frames": 33, "video_codec_options": {"preset": "ultrafast"}}' \ -o output.mp4 ``` +Set `preencode_batch_frames` in `extra_params` (or `extra_args` for offline +sampling) to tune batching. The worker accumulates complete VAE chunks until +it has at least this many frames, then transfers and encodes them together. +It always flushes the final partial batch. This is a threshold, not an exact +chunk length: a value of 1 submits every native chunk immediately, and a value +smaller than a native chunk does not split it. Larger values reduce transfers +but retain more frames on the accelerator and delay encoding. The VAE decode +window and output frame count stay unchanged. Wan S2V keeps its existing +per-clip behavior by default. Zero, negative, fractional, boolean, string, and +null values are rejected when pre-encoding is enabled. + `preencode_mp4` applies to the complete-MP4 response paths only. The `/v1/realtime/video` WebSocket endpoint rejects it, because that path already -overlaps encoding through its own incremental fragmented-MP4 encoder. +overlaps encoding through its own incremental fragmented-MP4 encoder. Wan also +rejects it together with `enable_frame_interpolation`, which needs the decoded +frames the pre-encoded path no longer materializes. -Support is per model: MiniMax-H3 implements it, and other models ignore the flag -and take the full-decode path. +Support is per model: MiniMax-H3 and Wan 2.2 (T2V, I2V, and S2V) implement it, +and other models ignore the flag and take the full-decode path. ## Storage diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index cea4862d1..7d2d8715f 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -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), diff --git a/tests/diffusion/models/wan2_2/test_wan22_chunked_mp4.py b/tests/diffusion/models/wan2_2/test_wan22_chunked_mp4.py new file mode 100644 index 000000000..a919d1727 --- /dev/null +++ b/tests/diffusion/models/wan2_2/test_wan22_chunked_mp4.py @@ -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 diff --git a/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py index 94f8b86fd..ab8cd132c 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py +++ b/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py @@ -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) diff --git a/tests/diffusion/utils/test_chunked_video.py b/tests/diffusion/utils/test_chunked_video.py new file mode 100644 index 000000000..849b49290 --- /dev/null +++ b/tests/diffusion/utils/test_chunked_video.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import io + +import av +import numpy as np +import pytest +import torch + +from vllm_omni.diffusion.utils.chunked_video import ( + ChunkedVideoMP4Session, + chunk_to_uint8_frames, + decode_to_mp4, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +class _FakeVAE: + """Publish committed chunks the way a chunked-decode VAE does.""" + + def __init__(self, value_range, *, batch=1, chunks=4, frames=2, height=16, width=32): + self.chunk_value_range = value_range + self._shape = (batch, 3, frames, height, width) + self._chunks = chunks + + def decode_with_chunks(self, z, *, on_chunk): + del z + low, high = self.chunk_value_range + for index in range(self._chunks): + value = low + (high - low) * index / max(self._chunks - 1, 1) + on_chunk(torch.full(self._shape, value)) + + +def _decoded(data: bytes): + with av.open(io.BytesIO(data)) as container: + return np.stack([f.to_ndarray(format="rgb24") for f in container.decode(video=0)]) + + +def test_quantization_maps_each_producer_range_onto_the_same_pixels(): + """A model's published range is what makes chunks comparable across models.""" + wan = torch.tensor([-1.0, 0.0, 1.0]).view(1, 1, 3, 1, 1).expand(1, 3, 3, 1, 1) + h3 = torch.tensor([0.0, 0.5, 1.0]).view(1, 1, 3, 1, 1).expand(1, 3, 3, 1, 1) + + from_wan = chunk_to_uint8_frames(wan.contiguous(), (-1.0, 1.0)) + from_h3 = chunk_to_uint8_frames(h3.contiguous(), (0.0, 1.0)) + + # Same physical pixels, different published ranges. + assert np.array_equal(from_wan, from_h3) + assert from_wan.reshape(-1, 3)[:, 0].tolist() == [0, 128, 255] + + +def test_quantization_rejects_a_degenerate_range(): + with pytest.raises(ValueError, match="increasing"): + chunk_to_uint8_frames(torch.zeros(1, 3, 1, 2, 2), (1.0, 1.0)) + + +@pytest.mark.parametrize("value_range", [(-1.0, 1.0), (0.0, 1.0)]) +def test_decode_to_mp4_drives_any_declared_range(value_range): + videos = decode_to_mp4(_FakeVAE(value_range), torch.zeros(1), fps=24, batch_frames=4) + + assert len(videos) == 1 + assert _decoded(videos[0]).shape[0] == 8 + + +def test_decode_to_mp4_emits_one_container_per_batch_entry(): + videos = decode_to_mp4(_FakeVAE((-1.0, 1.0), batch=3), torch.zeros(1), fps=24) + + assert len(videos) == 3 + assert all(_decoded(v).shape[0] == 8 for v in videos) + + +def test_decode_to_mp4_rejects_a_vae_without_the_capability(): + class PlainVAE: + def decode(self, z): + del z + + with pytest.raises(TypeError, match="chunked VAE decode capability"): + decode_to_mp4(PlainVAE(), torch.zeros(1), fps=24) + + +def test_decode_to_mp4_returns_nothing_for_a_rank_without_output(): + class SilentVAE: + chunk_value_range = (-1.0, 1.0) + + def decode_with_chunks(self, z, *, on_chunk): + del z, on_chunk + + assert decode_to_mp4(SilentVAE(), torch.zeros(1), fps=24) == [] + + +def test_session_crops_the_decoder_padding_away(): + session = ChunkedVideoMP4Session(value_range=(0.0, 1.0), fps=24, crop=(8, 16)) + session.push(torch.zeros(1, 3, 4, 16, 32)) + + assert _decoded(session.finish()[0]).shape[1:3] == (8, 16) + + +def test_session_muxes_one_waveform_per_batch_entry(): + audio = [np.zeros(8000, dtype=np.float32), np.full(8000, 0.5, dtype=np.float32)] + session = ChunkedVideoMP4Session(value_range=(0.0, 1.0), fps=24, audio_waveforms=audio, audio_sample_rate=16000) + session.push(torch.zeros(2, 3, 4, 16, 16)) + + for data in session.finish(): + with av.open(io.BytesIO(data)) as container: + assert container.streams.audio + + +def test_session_rejects_a_waveform_count_that_does_not_match_the_batch(): + session = ChunkedVideoMP4Session(value_range=(0.0, 1.0), fps=24, audio_waveforms=[np.zeros(8000, dtype=np.float32)]) + with pytest.raises(ValueError, match="one audio waveform per batch entry"): + session.push(torch.zeros(2, 3, 4, 16, 16)) + + +def test_session_rejects_a_non_positive_batch(): + with pytest.raises(ValueError, match="batch_frames"): + ChunkedVideoMP4Session(value_range=(0.0, 1.0), fps=24, batch_frames=0) diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index 60d6f4125..1a8395e5c 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -1726,66 +1726,47 @@ class MiniMaxH3Pipeline( ) -> bytes: """Decode and encode one output on the worker without full-video materialization. - Audio is decoded first so the incremental mux session can attach its audio - stream before temporal video chunks arrive. The callback receives committed - float32 ``BCTHW`` frames, performs the requested-size crop and uint8 - conversion in the worker, then applies bounded backpressure to the encoder. - """ - from vllm_omni.diffusion.utils.media_utils import ChunkedMP4Encoder + Audio is decoded first so the incremental mux session can attach its + audio stream before temporal video chunks arrive. Everything after a + chunk is committed -- crop, quantization, transfer, encoding -- is the + shared consumer's job; this method only supplies what is specific to + H3: the audio waveform, the requested-size crop, and the fixed rate. - if batch_frames <= 0: - raise ValueError("batch_frames must be positive") + Every rank of a distributed VAE group drives the temporal collectives, + but only the output owner receives chunks, so a peer rank returns empty + bytes -- the pre-encoded counterpart of the empty tensor the full + decode leaves there. + """ + from vllm_omni.diffusion.utils.chunked_video import decode_to_mp4 as decode_chunks_to_mp4 with self._component_on_device(self.audio_vae): audio = self.audio_vae.decode_latent(audio_latent) audio_np = audio.detach().float().cpu().numpy() if audio_np.ndim == 3 and audio_np.shape[0] == 1: audio_np = audio_np[0] - encoder = ChunkedMP4Encoder( - width=width, - height=height, - fps=MINIMAX_H3_FPS, - audio_waveform=audio_np, - audio_sample_rate=MINIMAX_H3_AUDIO_SAMPLE_RATE, - max_pending=max_pending, - video_codec_options=video_codec_options, - ) - pending_chunks: list[torch.Tensor] = [] - pending_frames = 0 - - def flush_pending() -> None: - nonlocal pending_frames - if not pending_chunks: - return - batched = torch.cat(pending_chunks, dim=1) - encoder.push(batched[0].cpu().numpy()) - pending_chunks.clear() - pending_frames = 0 - - def on_chunk(frames: torch.Tensor) -> None: - nonlocal pending_frames - prepared = _prepare_minimax_h3_video_output(frames[..., :height, :width]) - if prepared.shape[0] != 1: - raise ValueError("MiniMax H3 chunked MP4 encoding currently expects one output per decoder") - pending_chunks.append(prepared) - pending_frames += int(prepared.shape[1]) - if pending_frames >= batch_frames: - flush_pending() - - try: - with self._component_on_device(self.video_vae): - with current_omni_platform.create_autocast_context( - device_type=self.device.type, - dtype=torch.float16, - enabled=True, - ): - self.video_vae.decode_with_chunks(video_latent, on_chunk=on_chunk) - flush_pending() - return encoder.finish() - except BaseException: - encoder.abort() - raise + with self._component_on_device(self.video_vae): + with current_omni_platform.create_autocast_context( + device_type=self.device.type, + dtype=torch.float16, + enabled=True, + ): + videos = decode_chunks_to_mp4( + self.video_vae, + video_latent, + fps=MINIMAX_H3_FPS, + audio_waveforms=[audio_np], + audio_sample_rate=MINIMAX_H3_AUDIO_SAMPLE_RATE, + batch_frames=batch_frames, + max_pending=max_pending, + video_codec_options=video_codec_options, + crop=(height, width), + ) + if not videos: + return b"" + if len(videos) != 1: + raise ValueError("MiniMax H3 chunked MP4 encoding currently expects one output per decoder") + return videos[0] def decode( self, diff --git a/vllm_omni/diffusion/models/wan2_2/chunked_mp4.py b/vllm_omni/diffusion/models/wan2_2/chunked_mp4.py new file mode 100644 index 000000000..5b16a7a13 --- /dev/null +++ b/vllm_omni/diffusion/models/wan2_2/chunked_mp4.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Worker-side MP4 encoding for Wan VAE temporal chunks.""" + +from __future__ import annotations + +from typing import Any + +from vllm_omni.diffusion.utils.media_utils import normalize_preencode_batch_frames, normalize_video_codec_options + +# Matches the fallback the video serving layer applies when a request omits fps, +# so the bytes the worker encodes carry the rate the response advertises. +WAN_DEFAULT_OUTPUT_FPS = 24 + +# Keep worker-side encoding consistent with the one-shot serving encoder when +# a request does not choose codec options explicitly. +WAN_DEFAULT_VIDEO_CODEC_OPTIONS = {"preset": "ultrafast", "threads": "0"} + +# Wan publishes one latent frame group per callback, which is far finer than one +# host transfer is worth. 17 frames matches MiniMax-H3's native clip size. +WAN_DEFAULT_BATCH_FRAMES = 17 + + +def wan_preencoded_mp4_payload(video: Any) -> dict[str, Any] | None: + """Return a post-process payload when the worker already encoded the MP4. + + Returns ``None`` for the ordinary frame-tensor output so callers fall + through to their existing post-processing. + + No fps is reported here: :func:`resolve_wan_output_fps` mirrors the fallback + the serving layer applies, so the rate the response advertises already + matches the rate these bytes were encoded at. + """ + if isinstance(video, (bytes, bytearray, memoryview)): + video = [video] + if not isinstance(video, list) or not all(isinstance(item, (bytes, bytearray, memoryview)) for item in video): + return None + return {"payload": {"video": [bytes(item) for item in video]}, "metadata": {}} + + +def resolve_wan_preencode_mp4(sampling_params: Any, *, output_type: str) -> bool: + """Return whether this request wants the worker to pre-encode the MP4. + + Reject rather than silently downgrade the combinations the pre-encoded path + cannot serve: it emits MP4 bytes, so nothing downstream can interpolate + frames or hand back an array. + """ + extra_args = getattr(sampling_params, "extra_args", None) or {} + if not extra_args.get("preencode_mp4", False): + return False + if getattr(sampling_params, "enable_frame_interpolation", False): + raise ValueError( + "preencode_mp4 cannot be combined with enable_frame_interpolation: interpolation " + "needs the decoded frames that the pre-encoded path no longer materializes" + ) + if output_type != "np": + raise ValueError(f"preencode_mp4 returns MP4 bytes and cannot serve output_type={output_type!r}") + return True + + +def resolve_wan_preencode_batch_frames(sampling_params: Any, *, default: int = WAN_DEFAULT_BATCH_FRAMES) -> int: + """Read the MP4 batching threshold before starting expensive generation.""" + extra_args = getattr(sampling_params, "extra_args", None) or {} + return normalize_preencode_batch_frames(extra_args.get("preencode_batch_frames", default)) + + +def resolve_wan_video_codec_options(sampling_params: Any) -> dict[str, str] | None: + """Read the request's encoder options, preserving serving defaults. + + An absent key inherits the one-shot serving encoder's defaults. Explicit + ``None`` and an empty mapping remain distinct caller choices. + """ + extra_args = getattr(sampling_params, "extra_args", None) or {} + return normalize_video_codec_options(extra_args.get("video_codec_options", WAN_DEFAULT_VIDEO_CODEC_OPTIONS)) + + +def resolve_wan_output_fps(sampling_params: Any) -> int: + """Resolve the fps the worker encodes with. + + The serving layer resolves output fps from the request's ``fps`` field, not + from ``frame_rate``, so pre-encoding has to read the same field or the bytes + would carry a rate the response does not advertise. + """ + fps = getattr(sampling_params, "fps", None) + if isinstance(fps, list): + fps = fps[0] if fps else None + return int(fps) if fps else WAN_DEFAULT_OUTPUT_FPS diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index 915fb45f7..a971806b0 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -40,12 +40,20 @@ from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery 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.chunked_mp4 import ( + 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.models.wan2_2.scheduling_wan_euler import WanEulerScheduler from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanSelfAttention, WanTransformer3DModel from vllm_omni.diffusion.offloader import OffloadPlan 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.diffusion.utils.chunked_video import decode_to_mp4 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt from vllm_omni.platforms import current_omni_platform @@ -269,6 +277,9 @@ def get_wan22_post_process_func( output_type = sampling_params.output_type if output_type == "latent": return video + encoded = wan_preencoded_mp4_payload(video) + if encoded is not None: + return encoded video_metadata = {} if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False): video, multiplier = interpolate_video_tensor( @@ -711,6 +722,8 @@ class Wan22Pipeline( num_steps = 40 if common.num_inference_steps is None else common.num_inference_steps output_type = common.output_type or "np" + preencode_mp4 = resolve_wan_preencode_mp4(common, output_type=output_type) + preencode_batch_frames = resolve_wan_preencode_batch_frames(common) if preencode_mp4 else 17 num_outputs_per_prompt = common.num_outputs_per_prompt or 1 attention_kwargs: dict | None = None @@ -965,27 +978,36 @@ class Wan22Pipeline( latents.device, latents.dtype ) latents = latents / latents_std + latents_mean - decoded = self.vae.decode(latents, return_dict=False)[0] - # Distributed VAE decode uses broadcast_result=False, so only the - # output-owning rank receives the full [B, C, T, H, W] video; other - # ranks get an empty placeholder. Emit typed media only from the - # owning rank and keep the placeholder on the legacy output field, so - # the media batch-dimension check in split_diffusion_output_by_request - # does not trip on every non-owner rank. - if decoded.dim() == 5: - output = None - media = DiffusionMediaOutput( - video=VideoMediaOutput( - tensor=decoded, - spec=VideoTensorSpec( - layout=VideoTensorLayout.BCTHW, - encoding=VideoTensorEncoding.NORMALIZED_FLOAT, - value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, - ), - ) + if preencode_mp4: + output = decode_to_mp4( + self.vae, + latents, + fps=resolve_wan_output_fps(common), + batch_frames=preencode_batch_frames, + video_codec_options=resolve_wan_video_codec_options(common), ) else: - output = decoded + decoded = self.vae.decode(latents, return_dict=False)[0] + # Distributed VAE decode uses broadcast_result=False, so only the + # output-owning rank receives the full [B, C, T, H, W] video; other + # ranks get an empty placeholder. Emit typed media only from the + # owning rank and keep the placeholder on the legacy output field, so + # the media batch-dimension check in split_diffusion_output_by_request + # does not trip on every non-owner rank. + if decoded.dim() == 5: + output = None + media = DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=decoded, + spec=VideoTensorSpec( + layout=VideoTensorLayout.BCTHW, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ), + ) + ) + else: + output = decoded if DEBUG_PERF: current_omni_platform.synchronize() diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py index 1a75c40da..9f03f5853 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py @@ -31,6 +31,13 @@ from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero from vllm_omni.diffusion.models.utils import _load_json +from vllm_omni.diffusion.models.wan2_2.chunked_mp4 import ( + 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.models.wan2_2.pipeline_wan2_2 import ( _WAN_TEXT_ENCODER_OFFLOAD_PLAN, build_wan_scheduler, @@ -46,6 +53,7 @@ from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3 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.diffusion.utils.chunked_video import decode_to_mp4 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.platforms import current_omni_platform @@ -70,6 +78,9 @@ def get_wan22_i2v_post_process_func( output_type = sampling_params.output_type if output_type == "latent": return video + encoded = wan_preencoded_mp4_payload(video) + if encoded is not None: + return encoded video_metadata = {} if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False): video, multiplier = interpolate_video_tensor( @@ -507,6 +518,8 @@ class Wan22I2VPipeline( num_steps = 40 if common.num_inference_steps is None else common.num_inference_steps output_type = common.output_type or "np" + preencode_mp4 = resolve_wan_preencode_mp4(common, output_type=output_type) + preencode_batch_frames = resolve_wan_preencode_batch_frames(common) if preencode_mp4 else 17 num_outputs_per_prompt = common.num_outputs_per_prompt or 1 attention_kwargs: dict | None = None @@ -736,7 +749,16 @@ class Wan22I2VPipeline( latents.device, latents.dtype ) latents = latents / latents_std + latents_mean - output = self.vae.decode(latents, return_dict=False)[0] + if preencode_mp4: + output = decode_to_mp4( + self.vae, + latents, + fps=resolve_wan_output_fps(common), + batch_frames=preencode_batch_frames, + video_codec_options=resolve_wan_video_codec_options(common), + ) + else: + output = self.vae.decode(latents, return_dict=False)[0] if DEBUG_PERF: current_omni_platform.synchronize() diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py index 1d28b425d..3cbd727a7 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py @@ -33,6 +33,11 @@ from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineL from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportImageInput from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler +from vllm_omni.diffusion.models.wan2_2.chunked_mp4 import ( + resolve_wan_preencode_batch_frames, + resolve_wan_preencode_mp4, + resolve_wan_video_codec_options, +) from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( load_transformer_config, load_wan_weights_with_optional_gate, @@ -44,6 +49,7 @@ from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import ( from vllm_omni.diffusion.offloader.config import DIT_COMPONENT, selected_offload_components from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.utils.chunked_video import ChunkedVideoMP4Session from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.platforms import current_omni_platform @@ -63,6 +69,10 @@ _S2V_DEFAULT_NEG_PROMPT = ( _AUDIO_VIDEO_RATE = 30 +# S2V generates at a fixed rate; the worker-side encoder and the post-process +# payload must report the same one. +S2V_OUTPUT_FPS = 16 + def _s2v_audio_condition_key(audio: str | np.ndarray | None) -> tuple[Any, ...]: if isinstance(audio, str): @@ -280,6 +290,10 @@ def get_wan22_s2v_post_process_func( if output_type == "latent": return video + if isinstance(video, list) and all(isinstance(item, (bytes, bytearray, memoryview)) for item in video): + # The worker already muxed the audio into these bytes. + return {"video": [bytes(item) for item in video], "audio": None, "fps": S2V_OUTPUT_FPS} + processed_video = video_processor.postprocess_video(video, output_type=output_type) if audio_waveform is not None: @@ -287,7 +301,7 @@ def get_wan22_s2v_post_process_func( "video": processed_video, "audio": audio_waveform, "audio_sample_rate": audio_sr, - "fps": 16, + "fps": S2V_OUTPUT_FPS, } return processed_video @@ -1177,6 +1191,13 @@ class Wan22S2VPipeline( ) num_steps = common.num_inference_steps or num_inference_steps num_outputs_per_prompt = common.num_outputs_per_prompt or 1 + # The request's output_type is what serving honours downstream; the keyword + # argument only carries a default for direct callers. Reading it here is what + # lets pre-encoding reject pil, pt, and latent instead of returning MP4 bytes + # to a caller that asked for frames. + output_type = common.output_type or output_type or "np" + preencode_mp4 = resolve_wan_preencode_mp4(common, output_type=output_type) + preencode_batch_frames = resolve_wan_preencode_batch_frames(common, default=1) if preencode_mp4 else 1 if common.guidance_scale_provided: guidance_scale = common.guidance_scale @@ -1315,170 +1336,207 @@ class Wan22S2VPipeline( # ---- 6. Multi-clip autoregressive denoising ---- clips: list[torch.Tensor] = [] + # Encode each finished clip while the next one denoises, instead of + # concatenating the whole video first. + mp4_session = ( + ChunkedVideoMP4Session( + value_range=self.vae.chunk_value_range, + audio_waveforms=[raw_audio_waveforms[index // num_outputs_per_prompt] for index in range(batch_size)], + audio_sample_rate=raw_audio_sr, + fps=S2V_OUTPUT_FPS, + batch_frames=preencode_batch_frames, + video_codec_options=resolve_wan_video_codec_options(common), + ) + if preencode_mp4 + else None + ) # Keep a pixel-space buffer of the trailing motion_frames for the # autoregressive connection between clips. videos_last_frames = torch.zeros([batch_size, 3, motion_frames, height, width], dtype=dtype, device=device) - for r in range(num_repeat): - clip_generators = _make_clip_generators(seeds, generators, r, device) - # -- Noise -- - if r == 0 and request_latents is not None: - latents = request_latents.to(device=device, dtype=dtype) - else: - latents = torch.stack( - [ - self.prepare_latents( - infer_frames=infer_frames, - height=height, - width=width, - motion_frames=motion_frames, - lat_motion_frames=lat_motion_frames, - dtype=dtype, - device=device, - generator=clip_generator, - ) - for clip_generator in clip_generators - ] + try: + for r in range(num_repeat): + clip_generators = _make_clip_generators(seeds, generators, r, device) + # -- Noise -- + if r == 0 and request_latents is not None: + latents = request_latents.to(device=device, dtype=dtype) + else: + latents = torch.stack( + [ + self.prepare_latents( + infer_frames=infer_frames, + height=height, + width=width, + motion_frames=motion_frames, + lat_motion_frames=lat_motion_frames, + dtype=dtype, + device=device, + generator=clip_generator, + ) + for clip_generator in clip_generators + ] + ) + + # -- Scheduler -- + self.scheduler.set_timesteps(num_steps, device=device, shift=self._flow_shift) + timesteps = self.scheduler.timesteps + self._num_timesteps = len(timesteps) + + # -- Slice audio for this clip -- + left_idx = r * infer_frames + right_idx = left_idx + infer_frames + audio_input = audio_emb[..., left_idx:right_idx] + + # -- Pose condition for this clip -- + # Default: zero condition (no pose driving) + # Shape: [B, C=16, T, H, W] + lat_h = height // self.vae_scale_factor_spatial + lat_w = width // self.vae_scale_factor_spatial + lat_target_frames = (infer_frames + 3 + motion_frames) // 4 - lat_motion_frames + cond_latents = torch.zeros( + [batch_size, 16, lat_target_frames, lat_h, lat_w], dtype=dtype, device=device ) - # -- Scheduler -- - self.scheduler.set_timesteps(num_steps, device=device, shift=self._flow_shift) - timesteps = self.scheduler.timesteps - self._num_timesteps = len(timesteps) + # -- Clone motion latents for this clip -- + input_motion_latents = motion_latents.clone() - # -- Slice audio for this clip -- - left_idx = r * infer_frames - right_idx = left_idx + infer_frames - audio_input = audio_emb[..., left_idx:right_idx] + # Max sequence length for the transformer + max_seq_len = int(np.prod([lat_target_frames, lat_h, lat_w]) // 4) - # -- Pose condition for this clip -- - # Default: zero condition (no pose driving) - # Shape: [B, C=16, T, H, W] - lat_h = height // self.vae_scale_factor_spatial - lat_w = width // self.vae_scale_factor_spatial - lat_target_frames = (infer_frames + 3 + motion_frames) // 4 - lat_motion_frames - cond_latents = torch.zeros([batch_size, 16, lat_target_frames, lat_h, lat_w], dtype=dtype, device=device) - - # -- Clone motion latents for this clip -- - input_motion_latents = motion_latents.clone() - - # Max sequence length for the transformer - max_seq_len = int(np.prod([lat_target_frames, lat_h, lat_w]) // 4) - - # -- Precompute audio embeddings once per clip -- - mf = [motion_frames, lat_motion_frames] - # When CPU offload is active the transformer lives on CPU until - # its __call__ hook fires. encode_audio() bypasses __call__, - # so move the audio sub-module to GPU explicitly. - _audio_enc = getattr(self.transformer, "casual_audio_encoder", None) - _audio_on_cpu = _audio_enc is not None and next(_audio_enc.parameters()).device.type == "cpu" - if _audio_on_cpu: - _audio_enc.to(device) - param_dtype = next(self.transformer.parameters()).dtype - with torch.amp.autocast(device.type, dtype=param_dtype): - positive_audio_emb = self.transformer.encode_audio(audio_input, mf) - do_true_cfg = self.do_classifier_free_guidance and negative_prompt_embeds is not None - if do_true_cfg: + # -- Precompute audio embeddings once per clip -- + mf = [motion_frames, lat_motion_frames] + # When CPU offload is active the transformer lives on CPU until + # its __call__ hook fires. encode_audio() bypasses __call__, + # so move the audio sub-module to GPU explicitly. + _audio_enc = getattr(self.transformer, "casual_audio_encoder", None) + _audio_on_cpu = _audio_enc is not None and next(_audio_enc.parameters()).device.type == "cpu" + if _audio_on_cpu: + _audio_enc.to(device) + param_dtype = next(self.transformer.parameters()).dtype with torch.amp.autocast(device.type, dtype=param_dtype): - negative_audio_emb = self.transformer.encode_audio(0.0 * audio_input, mf) - else: - negative_audio_emb = None - if _audio_on_cpu: - _audio_enc.to("cpu") + positive_audio_emb = self.transformer.encode_audio(audio_input, mf) + do_true_cfg = self.do_classifier_free_guidance and negative_prompt_embeds is not None + if do_true_cfg: + with torch.amp.autocast(device.type, dtype=param_dtype): + negative_audio_emb = self.transformer.encode_audio(0.0 * audio_input, mf) + else: + negative_audio_emb = None + if _audio_on_cpu: + _audio_enc.to("cpu") - # -- Denoising loop -- - latents = self.diffuse( - latents=latents, - timesteps=timesteps, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - guidance_scale=guidance_scale, - clip_generator=clip_generators, - dtype=dtype, - device=device, - max_seq_len=max_seq_len, - cond_latents=cond_latents, - input_motion_latents=input_motion_latents, - ref_latents=ref_latents, - motion_frames=mf, - drop_first_motion=drop_first_motion and r == 0, - positive_audio_emb=positive_audio_emb, - negative_audio_emb=negative_audio_emb, - ) - - # ---- Decode this clip ---- - if self._should_release_dit_before_decode(): - self.transformer.to("cpu") - current_omni_platform.empty_cache() - - if not (drop_first_motion and r == 0): - decode_latents = torch.cat([motion_latents, latents], dim=2) - else: - decode_latents = torch.cat([ref_latents, latents], dim=2) - - decode_latents = self._denormalize_latents(decode_latents) - decode_latents = decode_latents.to(self.vae.dtype) - clip_video = self.vae.decode(decode_latents, return_dict=False)[0] # [1, C, T, H, W] - - # Handle VAE patch parallel: only rank 0 gets result, broadcast to all ranks - # This is needed for S2V's autoregressive loop where all ranks need the decoded frames - if clip_video.numel() == 0: - # Non-rank0 received empty tensor from patch parallel decode - # Use the same broadcast mechanism as the VAE patch parallel code - import torch.distributed as dist - - total_frames = decode_latents.shape[2] - - # Create buffer for broadcast - clip_video = torch.empty( - (batch_size, 3, total_frames, height, width), - device=decode_latents.device, - dtype=decode_latents.dtype, + # -- Denoising loop -- + latents = self.diffuse( + latents=latents, + timesteps=timesteps, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + guidance_scale=guidance_scale, + clip_generator=clip_generators, + dtype=dtype, + device=device, + max_seq_len=max_seq_len, + cond_latents=cond_latents, + input_motion_latents=input_motion_latents, + ref_latents=ref_latents, + motion_frames=mf, + drop_first_motion=drop_first_motion and r == 0, + positive_audio_emb=positive_audio_emb, + negative_audio_emb=negative_audio_emb, ) - # Get the VAE's patch parallel group (same one used in decode) - vae_pp_group = getattr(self.vae, "_vae_pp_group", None) - if vae_pp_group is not None: - # Broadcast using the same group as VAE patch parallel - dist.broadcast(clip_video, src=0, group=vae_pp_group) + # ---- Decode this clip ---- + if self._should_release_dit_before_decode(): + self.transformer.to("cpu") + current_omni_platform.empty_cache() - # Trim to the infer_frames of interest - clip_video = clip_video[:, :, -infer_frames:] - if drop_first_motion and r == 0: - # Drop the first 3 frames (artifact from ref_latents prepend) - clip_video = clip_video[:, :, 3:] + if not (drop_first_motion and r == 0): + decode_latents = torch.cat([motion_latents, latents], dim=2) + else: + decode_latents = torch.cat([ref_latents, latents], dim=2) - # ---- Update motion for next clip (autoregressive) ---- - overlap_frames_num = min(motion_frames, clip_video.shape[2]) - videos_last_frames = torch.cat( - [videos_last_frames[:, :, overlap_frames_num:], clip_video[:, :, -overlap_frames_num:]], - dim=2, + decode_latents = self._denormalize_latents(decode_latents) + decode_latents = decode_latents.to(self.vae.dtype) + clip_video = self.vae.decode(decode_latents, return_dict=False)[0] # [1, C, T, H, W] + # Under VAE patch parallelism only the output-owning rank decodes a + # real clip; peers get an empty placeholder filled by the broadcast. + owns_output = clip_video.numel() > 0 + + # Handle VAE patch parallel: only rank 0 gets result, broadcast to all ranks + # This is needed for S2V's autoregressive loop where all ranks need the decoded frames + if clip_video.numel() == 0: + # Non-rank0 received empty tensor from patch parallel decode + # Use the same broadcast mechanism as the VAE patch parallel code + import torch.distributed as dist + + total_frames = decode_latents.shape[2] + + # Create buffer for broadcast + clip_video = torch.empty( + (batch_size, 3, total_frames, height, width), + device=decode_latents.device, + dtype=decode_latents.dtype, + ) + + # Get the VAE's patch parallel group (same one used in decode) + vae_pp_group = getattr(self.vae, "_vae_pp_group", None) + if vae_pp_group is not None: + # Broadcast using the same group as VAE patch parallel + dist.broadcast(clip_video, src=0, group=vae_pp_group) + + # Trim to the infer_frames of interest + clip_video = clip_video[:, :, -infer_frames:] + if drop_first_motion and r == 0: + # Drop the first 3 frames (artifact from ref_latents prepend) + clip_video = clip_video[:, :, 3:] + + # ---- Update motion for next clip (autoregressive) ---- + overlap_frames_num = min(motion_frames, clip_video.shape[2]) + videos_last_frames = torch.cat( + [videos_last_frames[:, :, overlap_frames_num:], clip_video[:, :, -overlap_frames_num:]], + dim=2, + ) + videos_last_frames = videos_last_frames.to(dtype=dtype, device=device) + # Only prepare motion latents if there's another clip coming + if r < num_repeat - 1: + motion_latents = self.prepare_motion_latents(videos_last_frames, device=device).to(dtype=dtype) + + if mp4_session is not None: + # Every rank needs the frames for the motion loop above, but the + # executor keeps only the owner's response, so a peer skips the + # encode and finishes with no containers, like T2V's non-owners. + if owns_output: + mp4_session.push(clip_video) + else: + clips.append(clip_video.cpu()) + + # Free VRAM between clips + if current_omni_platform.is_available(): + current_omni_platform.empty_cache() + + # ---- Assemble the output ---- + if mp4_session is not None: + output = mp4_session.finish() # one progressive MP4 per batch entry + else: + output = torch.cat(clips, dim=2) # [B, C, T_total, H, W] + + outputs = split_diffusion_output_by_request( + DiffusionOutput( + output=output, + stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, + ), + req, + num_outputs_per_prompt=num_outputs_per_prompt, ) - videos_last_frames = videos_last_frames.to(dtype=dtype, device=device) - # Only prepare motion latents if there's another clip coming - if r < num_repeat - 1: - motion_latents = self.prepare_motion_latents(videos_last_frames, device=device).to(dtype=dtype) - - clips.append(clip_video.cpu()) - - # Free VRAM between clips - if current_omni_platform.is_available(): - current_omni_platform.empty_cache() - - # ---- Concatenate all clips ---- - output = torch.cat(clips, dim=2) # [B, C, T_total, H, W] - - outputs = split_diffusion_output_by_request( - DiffusionOutput( - output=output, - stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, - ), - req, - num_outputs_per_prompt=num_outputs_per_prompt, - ) - for request_output, raw_audio_waveform in zip(outputs, raw_audio_waveforms): - request_output.output = (request_output.output, raw_audio_waveform, raw_audio_sr) - return outputs + if mp4_session is not None: + # The waveform is already muxed into each container. + return outputs + for request_output, raw_audio_waveform in zip(outputs, raw_audio_waveforms): + request_output.output = (request_output.output, raw_audio_waveform, raw_audio_sr) + return outputs + except BaseException: + if mp4_session is not None: + mp4_session.abort() + raise def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return load_wan_weights_with_optional_gate(self, weights) diff --git a/vllm_omni/diffusion/utils/chunked_video.py b/vllm_omni/diffusion/utils/chunked_video.py new file mode 100644 index 000000000..79ceeab8c --- /dev/null +++ b/vllm_omni/diffusion/utils/chunked_video.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Turn committed VAE chunks into MP4 bytes without materializing the video. + +The producer side is the model's business: a Wan VAE walks a causal +frame-by-frame loop, a MiniMax-H3 VAE walks overlapping clips, and Wan S2V +decodes one clip per autoregressive iteration. What happens to a chunk once it +is final is not: every one of them quantizes to uint8, moves to the host once, +and lands in a bounded encoder. That consumer lives here so a model only has to +publish finished chunks -- :class:`SupportsChunkedVAEDecode` for producers that +can be driven, :class:`ChunkedVideoMP4Session` for producers that drive +themselves. +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import torch + +from vllm_omni.diffusion.models.interface import supports_chunked_vae_decode +from vllm_omni.diffusion.utils.media_utils import ChunkedMP4Encoder + + +def chunk_to_uint8_frames(chunk: torch.Tensor, value_range: tuple[float, float]) -> np.ndarray: + """Quantize a ``BCTHW`` chunk to ``BTHWC`` uint8, transferring once. + + ``value_range`` is the interval the producer publishes, which differs per + checkpoint (see ``SupportsChunkedVAEDecode.chunk_value_range``). Quantizing + on the accelerator means one transfer moves the final bytes rather than + float frames. + + That transfer is synchronous: ``.cpu()`` blocks inside the producer's chunk + callback, so what overlaps the remaining decode today is the CPU H.264 + encode of earlier chunks, not this copy. Splitting device preparation, + device-to-host copy, and encode into three independently scheduled stages -- + a pool of reusable pinned host slots with non-blocking copies on a dedicated + stream, leases released by the encoder -- fits behind this same callback + contract, but is not implemented here. + """ + low, high = float(value_range[0]), float(value_range[1]) + if high <= low: + raise ValueError(f"value_range must be increasing, got {value_range!r}") + scale = 255.0 / (high - low) + frames = chunk.clamp(low, high).sub(low).mul(scale).round().to(torch.uint8) + return frames.permute(0, 2, 3, 4, 1).cpu().numpy() + + +class ChunkedVideoMP4Session: + """Encode committed video chunks into one progressive MP4 per batch entry. + + Push finished ``BCTHW`` chunks as the producer commits them; each batch + entry gets its own bounded encoder, so host transfer and H.264 encoding + overlap whatever the producer is still decoding. + + ``audio_waveforms`` holds one waveform per batch entry (``None`` for a + silent entry), so a caller with several outputs per prompt repeats a + request's waveform across its entries. ``batch_frames`` coalesces transfers + for producers that publish finer than a transfer is worth; ``crop`` trims + the decoder's padding to the requested output size. + """ + + def __init__( + self, + *, + value_range: tuple[float, float], + fps: float, + audio_waveforms: list[np.ndarray | None] | None = None, + audio_sample_rate: int | None = None, + batch_frames: int = 1, + max_pending: int = 2, + video_codec_options: dict[str, str] | None = None, + crop: tuple[int, int] | None = None, + ) -> None: + if batch_frames <= 0: + raise ValueError("batch_frames must be positive") + self._value_range = value_range + self._fps = fps + self._audio_waveforms = audio_waveforms + self._audio_sample_rate = audio_sample_rate + self._batch_frames = batch_frames + self._max_pending = max_pending + self._video_codec_options = video_codec_options + self._crop = crop + self._encoders: list[ChunkedMP4Encoder] = [] + self._pending: list[torch.Tensor] = [] + self._pending_frames = 0 + + def push(self, chunk: torch.Tensor) -> None: + """Queue one committed ``BCTHW`` chunk.""" + if self._crop is not None: + height, width = self._crop + chunk = chunk[..., :height, :width] + self._pending.append(chunk) + self._pending_frames += int(chunk.shape[2]) + if self._pending_frames >= self._batch_frames: + self._flush() + + def finish(self) -> list[bytes]: + """Flush what is pending and return one MP4 per batch entry.""" + self._flush() + return [encoder.finish() for encoder in self._encoders] + + def abort(self) -> None: + for encoder in self._encoders: + encoder.abort() + + def _flush(self) -> None: + if not self._pending: + return + frames = chunk_to_uint8_frames(torch.cat(self._pending, dim=2), self._value_range) + if not self._encoders: + self._encoders = [ + ChunkedMP4Encoder( + width=frames.shape[3], + height=frames.shape[2], + fps=self._fps, + audio_waveform=self._waveform_for(index, frames.shape[0]), + audio_sample_rate=self._audio_sample_rate, + max_pending=self._max_pending, + video_codec_options=self._video_codec_options, + ) + for index in range(frames.shape[0]) + ] + for index, encoder in enumerate(self._encoders): + encoder.push(np.ascontiguousarray(frames[index])) + self._pending.clear() + self._pending_frames = 0 + + def _waveform_for(self, index: int, batch_size: int) -> np.ndarray | None: + if self._audio_waveforms is None: + return None + if len(self._audio_waveforms) != batch_size: + raise ValueError( + f"expected one audio waveform per batch entry, got " + f"{len(self._audio_waveforms)} for {batch_size} entries" + ) + return self._audio_waveforms[index] + + +def decode_to_mp4(vae: Any, z: torch.Tensor, **session_kwargs: Any) -> list[bytes]: + """Decode ``z`` straight into one progressive MP4 per batch entry. + + Drives a VAE that declares :class:`SupportsChunkedVAEDecode`, so host + transfer and encoding overlap the remaining decode and the full video is + never materialized. Ranks that own no decode output receive no chunks and + get an empty list, matching the empty tensor the full-decode path returns + there. + """ + if not supports_chunked_vae_decode(vae): + raise TypeError(f"{type(vae).__name__} does not expose the chunked VAE decode capability") + + session = ChunkedVideoMP4Session(value_range=vae.chunk_value_range, **session_kwargs) + try: + vae.decode_with_chunks(z, on_chunk=session.push) + return session.finish() + except BaseException: + session.abort() + raise diff --git a/vllm_omni/diffusion/utils/media_utils.py b/vllm_omni/diffusion/utils/media_utils.py index 99d20a0af..b2141400c 100644 --- a/vllm_omni/diffusion/utils/media_utils.py +++ b/vllm_omni/diffusion/utils/media_utils.py @@ -51,8 +51,11 @@ def _validate_video_chunk(chunk: np.ndarray, *, width: int, height: int) -> None class ChunkedMP4Encoder: """Encode temporal video chunks while the producer is still decoding. - A bounded queue and one muxing worker provide ordered backpressure while - keeping host memory bounded by ``max_pending`` chunks. Chunks use the same + A bounded queue and one muxing worker provide ordered backpressure. What + that bounds is pending raw-frame memory: at most ``max_pending`` chunks wait + to be encoded, so the bound is in chunks rather than bytes and a taller + frame or a longer chunk raises it. The muxed container is not bounded -- it + accumulates in memory and grows with output duration. Chunks use the same ``(T, H, W, 3)`` uint8 contract as :func:`mux_video_audio_bytes`. """