# SPDX-License-Identifier: Apache-2.0 # 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 import PIL.Image import pytest import torch import torch.nn as nn from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import ( Wan22S2VPipeline, _make_clip_generators, ) from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import WanS2VTransformer3DModel from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @pytest.mark.parametrize("configured_shift", [None, 12.0]) def test_s2v_constructor_preserves_unshifted_endpoints(configured_shift: float | None) -> None: module = "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v" config = SimpleNamespace(model="unused", flow_shift=configured_shift, enable_diffusion_pipeline_profiler=False) def init_components(pipeline, *args): pipeline.vae = SimpleNamespace(config=SimpleNamespace(scale_factor_temporal=4, scale_factor_spatial=8)) with ( patch(f"{module}.get_local_device", return_value=torch.device("cpu")), patch(f"{module}._resolve_model_path", return_value="unused"), patch(f"{module}._is_diffusers_format", return_value=True), patch.object(Wan22S2VPipeline, "_init_diffusers_format", init_components), patch.object(Wan22S2VPipeline, "setup_diffusion_pipeline_profiler"), ): pipeline = Wan22S2VPipeline(od_config=config) assert pipeline._flow_shift == (3.0 if configured_shift is None else configured_shift) assert pipeline.scheduler.config.shift == pipeline.scheduler.config["shift"] == 1.0 assert pipeline.scheduler.sigma_max == float(np.float32(0.999)) def _make_s2v_sampling(**overrides): values: dict[str, Any] = { "height": 16, "width": 16, "num_frames": 8, "num_inference_steps": 1, "guidance_scale_provided": True, "guidance_scale": 1.0, "generator": None, "seed": None, "num_outputs_per_prompt": 1, "max_sequence_length": 8, "latents": None, "output_type": "latent", "extra_args": {}, } values.update(overrides) return SimpleNamespace(**values) def _make_s2v_validation_pipeline() -> Wan22S2VPipeline: pipeline = object.__new__(Wan22S2VPipeline) nn.Module.__init__(pipeline) pipeline.device = torch.device("cpu") pipeline.transformer = SimpleNamespace(dtype=torch.float32) 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.check_inputs = lambda *args, **kwargs: None pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None) # type: ignore[method-assign] return pipeline def test_s2v_predict_noise_keeps_batch_dimension() -> None: pipeline = object.__new__(Wan22S2VPipeline) nn.Module.__init__(pipeline) pipeline.device = torch.device("cpu") transformer = MagicMock() transformer.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)]) transformer.return_value = (torch.zeros(1, 16, 2, 2, 2),) pipeline.transformer = transformer result = pipeline.predict_noise(hidden_states=torch.zeros(1, 16, 2, 2, 2)) assert result.shape == (1, 16, 2, 2, 2) def test_s2v_clip_generators_preserve_main_seed_per_clip_behavior() -> None: explicit_generator = torch.Generator(device="cpu").manual_seed(999) generators = _make_clip_generators( seeds=[1234, 5678, None, None], request_generators=[ torch.Generator(device="cpu").manual_seed(1234), torch.Generator(device="cpu").manual_seed(5678), explicit_generator, None, ], clip_index=2, device=torch.device("cpu"), ) assert [generator.initial_seed() for generator in generators if generator is not None] == [1236, 5680, 999, 2] assert generators[2] is explicit_generator def test_s2v_forward_rejects_different_num_repeat() -> None: pipeline = _make_s2v_validation_pipeline() pipeline.encode_audio = lambda *args, **kwargs: (torch.zeros(1, 1, 2, 16), 2, 16) # type: ignore[method-assign] image = PIL.Image.new("RGB", (16, 16)) audio = np.zeros(16000, dtype=np.float32) batch = DiffusionRequestBatch( requests=[ SimpleNamespace( request_id="a", prompt={ "prompt": "first", "multi_modal_data": {"image": image, "audio": audio}, "additional_information": {"num_repeat": 1}, }, sampling_params=_make_s2v_sampling(), ), SimpleNamespace( request_id="b", prompt={ "prompt": "second", "multi_modal_data": {"image": image, "audio": audio}, "additional_information": {"num_repeat": 2}, }, sampling_params=_make_s2v_sampling(), ), ] ) with pytest.raises(ValueError, match="num_repeat"): pipeline.forward(batch) def test_s2v_forward_rejects_different_audio_time_lengths() -> None: pipeline = _make_s2v_validation_pipeline() pipeline.encode_audio = lambda audio, **kwargs: ( # type: ignore[method-assign] torch.zeros(1, 1, 2, 16 if audio[0] == 0 else 24), 1, 16 if audio[0] == 0 else 24, ) image = PIL.Image.new("RGB", (16, 16)) batch = DiffusionRequestBatch( requests=[ SimpleNamespace( request_id="a", prompt={ "prompt": "first", "multi_modal_data": {"image": image, "audio": np.zeros(16000, dtype=np.float32)}, }, sampling_params=_make_s2v_sampling(), ), SimpleNamespace( request_id="b", prompt={ "prompt": "second", "multi_modal_data": {"image": image, "audio": np.ones(16000, dtype=np.float32)}, }, sampling_params=_make_s2v_sampling(), ), ] ) with pytest.raises(ValueError, match="audio embedding time lengths"): pipeline.forward(batch) def test_s2v_forward_rejects_different_init_first_frame() -> None: pipeline = _make_s2v_validation_pipeline() image = PIL.Image.new("RGB", (16, 16)) audio = np.zeros(16000, dtype=np.float32) batch = DiffusionRequestBatch( requests=[ SimpleNamespace( request_id="a", prompt={ "prompt": "first", "multi_modal_data": {"image": image, "audio": audio}, "additional_information": {"init_first_frame": True}, }, sampling_params=_make_s2v_sampling(), ), SimpleNamespace( request_id="b", prompt={ "prompt": "second", "multi_modal_data": {"image": image, "audio": audio}, "additional_information": {"init_first_frame": False}, }, sampling_params=_make_s2v_sampling(), ), ] ) with pytest.raises(ValueError, match="init_first_frame"): pipeline.forward(batch) def test_s2v_forward_rejects_different_raw_audio_shapes() -> None: pipeline = _make_s2v_validation_pipeline() pipeline.encode_audio = lambda *args, **kwargs: (torch.zeros(1, 1, 2, 16), 1, 16) # type: ignore[method-assign] image = PIL.Image.new("RGB", (16, 16)) batch = DiffusionRequestBatch( requests=[ SimpleNamespace( request_id="a", prompt={ "prompt": "first", "multi_modal_data": {"image": image, "audio": np.zeros(16000, dtype=np.float32)}, }, sampling_params=_make_s2v_sampling(), ), SimpleNamespace( request_id="b", prompt={ "prompt": "second", "multi_modal_data": {"image": image, "audio": np.zeros(8000, dtype=np.float32)}, }, sampling_params=_make_s2v_sampling(), ), ] ) with pytest.raises(ValueError, match="raw audio shapes and sample rates"): pipeline.forward(batch) @pytest.mark.parametrize("shift", [3.0, 12.0]) def test_s2v_forward_batches_request_local_inputs_and_splits_outputs(shift: float) -> None: pipeline = object.__new__(Wan22S2VPipeline) nn.Module.__init__(pipeline) pipeline.device = torch.device("cpu") pipeline.transformer = MagicMock() pipeline.transformer.dtype = torch.float32 pipeline.transformer.parameters.return_value = 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 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._flow_shift = shift pipeline.scheduler = FlowUniPCMultistepScheduler(shift=1.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"]) image = PIL.Image.new("RGB", (16, 16)) audio_a = np.zeros(16000, dtype=np.float32) audio_b = np.ones(16000, dtype=np.float32) latents_a = torch.zeros(2, 16, 2, 2, 2) latents_b = torch.ones(2, 16, 2, 2, 2) generators_a = [torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)] generators_b = [torch.Generator().manual_seed(3), torch.Generator().manual_seed(4)] batch = DiffusionRequestBatch( requests=[ SimpleNamespace( request_id="a", prompt={ "prompt": "first", "negative_prompt": "negative first", "multi_modal_data": {"image": image, "audio": audio_a}, }, sampling_params=_make_s2v_sampling( num_inference_steps=5, num_outputs_per_prompt=2, generator=generators_a, latents=latents_a, ), ), SimpleNamespace( request_id="b", prompt={ "prompt": "second", "negative_prompt": "negative second", "multi_modal_data": {"image": image, "audio": audio_b}, }, sampling_params=_make_s2v_sampling( num_inference_steps=5, num_outputs_per_prompt=2, generator=generators_b, latents=latents_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) expected = np.linspace(float(np.float32(0.999)), 0.0, 6)[:-1] expected = shift * expected / (1.0 + (shift - 1.0) * expected) torch.testing.assert_close( pipeline.scheduler.sigmas, torch.tensor(np.append(expected, 0.0), dtype=torch.float32), rtol=0, atol=0 ) assert pipeline.scheduler.config.shift == pipeline.scheduler.config["shift"] == 1.0 assert len(outputs) == 2 assert outputs[0].output[0].shape[0] == 2 assert outputs[1].output[0].shape[0] == 2 assert outputs[0].output[1].shape == (16000,) assert outputs[1].output[1].shape == (16000,) np.testing.assert_array_equal(outputs[0].output[1], audio_a) np.testing.assert_array_equal(outputs[1].output[1], audio_b) assert np.shares_memory(outputs[0].output[1], audio_a) assert np.shares_memory(outputs[1].output[1], audio_b) torch.testing.assert_close(pipeline.diffuse.call_args.kwargs["latents"], torch.cat([latents_a, latents_b])) assert pipeline.diffuse.call_args.kwargs["clip_generator"] == generators_a + generators_b assert pipeline.encode_prompt.call_args.kwargs["prompt"] == ["first", "second"] assert pipeline.encode_prompt.call_args.kwargs["negative_prompt"] == ["negative first", "negative second"] def test_s2v_exposes_hsdp_shard_conditions_for_transformer_blocks(): model = object.__new__(WanS2VTransformer3DModel) nn.Module.__init__(model) model.blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(3)]) conditions = getattr(model, "_hsdp_shard_conditions", None) assert conditions is not None assert len(conditions) == 1 matched = [] for name, module in model.named_modules(): if any(cond(name, module) for cond in conditions): matched.append(name) assert matched == ["blocks.0", "blocks.1", "blocks.2"] def test_s2v_hsdp_shard_condition_does_not_match_non_block_modules(): model = object.__new__(WanS2VTransformer3DModel) nn.Module.__init__(model) model.blocks = nn.ModuleList([nn.Linear(4, 4)]) model.head_indicator = nn.Linear(4, 4) model.casual_audio_encoder = nn.Linear(4, 4) conditions = model._hsdp_shard_conditions non_block_matched = [] for name, module in model.named_modules(): if name and "blocks" not in name: if any(cond(name, module) for cond in conditions): non_block_matched.append(name) assert non_block_matched == [] def test_encode_audio_calls_unshard_reshard_when_fsdp_managed(): model = object.__new__(WanS2VTransformer3DModel) nn.Module.__init__(model) model.enable_adain = False model.casual_audio_encoder = MagicMock(return_value=torch.zeros(1, 10, 64)) model.unshard = MagicMock() model.reshard = MagicMock() audio_input = torch.randn(1, 1, 64, 5) motion_frames = [2, 2] result = model.encode_audio(audio_input, motion_frames) model.unshard.assert_called_once() model.reshard.assert_called_once() assert "audio_emb" in result def test_encode_audio_reshard_called_on_exception(): """Test that reshard() is always called even when encode_audio logic raises.""" model = object.__new__(WanS2VTransformer3DModel) nn.Module.__init__(model) model.enable_adain = False model.casual_audio_encoder = MagicMock(side_effect=RuntimeError("encoder failed")) model.unshard = MagicMock() model.reshard = MagicMock() audio_input = torch.randn(1, 1, 64, 5) motion_frames = [2, 2] with pytest.raises(RuntimeError, match="encoder failed"): model.encode_audio(audio_input, motion_frames) model.unshard.assert_called_once() model.reshard.assert_called_once() def test_encode_audio_skips_unshard_reshard_when_not_fsdp(): model = object.__new__(WanS2VTransformer3DModel) nn.Module.__init__(model) model.enable_adain = False model.casual_audio_encoder = MagicMock(return_value=torch.zeros(1, 10, 64)) audio_input = torch.randn(1, 1, 64, 5) motion_frames = [2, 2] result = model.encode_audio(audio_input, motion_frames) assert not hasattr(model, "unshard") assert not hasattr(model, "reshard") assert "audio_emb" in result @pytest.mark.parametrize( ("components", "use_hsdp", "expected"), [ (None, False, True), (["dit"], False, True), (["text_encoder"], False, False), (["dit"], True, False), ], ) def test_s2v_dit_release_honors_component_selection(components, use_hsdp, expected): pipeline = object.__new__(Wan22S2VPipeline) nn.Module.__init__(pipeline) compact = None if components is None else {"mode": "module", "components": components} pipeline.od_config = SimpleNamespace( diffusion_offload_config=compact, enable_cpu_offload=True, enable_layerwise_offload=False, enable_distributed_layerwise_offload=False, dlo_use_allgather=True, dlo_resident_layers=0, pin_cpu_memory=True, parallel_config=SimpleNamespace(use_hsdp=use_hsdp), ) assert pipeline._should_release_dit_before_decode() is expected def test_s2v_pipeline_hsdp_forward_complete_process(): """Integration-level mock test verifying the complete forward path of the S2V pipeline under HSDP mode. Verifies: - Text encoding is invoked - Audio encoding invokes unshard/reshard on the transformer - Reference image encoding via VAE - Denoising loop calls the transformer - VAE decode is invoked - transformer.to('cpu') is NOT called (HSDP mode) - Final output is a DiffusionOutput with video + audio """ import numpy as np import PIL.Image from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import Wan22S2VPipeline pipeline = object.__new__(Wan22S2VPipeline) nn.Module.__init__(pipeline) # -- Config -- od_config = MagicMock() od_config.enable_cpu_offload = True od_config.enable_diffusion_pipeline_profiler = False parallel_config = MagicMock() parallel_config.use_hsdp = True od_config.parallel_config = parallel_config pipeline.od_config = od_config # -- Pipeline attributes -- pipeline.device = torch.device("cpu") pipeline._guidance_scale = 4.5 pipeline._num_timesteps = None pipeline._current_timestep = None pipeline.vae_scale_factor_spatial = 8 pipeline.vae_scale_factor_temporal = 4 pipeline.resolution_divisor = 16 pipeline.motion_frames = 73 pipeline.drop_first_motion = True pipeline.fps = 16 pipeline.audio_sample_m = 0 pipeline._DEFAULT_INFER_FRAMES = 80 # -- Mock text encoder -- pipeline.tokenizer = MagicMock() pipeline.tokenizer.return_value = MagicMock( input_ids=torch.zeros(1, 512, dtype=torch.long), attention_mask=torch.ones(1, 512, dtype=torch.long), ) mock_text_encoder = MagicMock() mock_text_encoder.dtype = torch.bfloat16 mock_text_encoder.return_value = MagicMock(last_hidden_state=torch.zeros(1, 512, 4096)) pipeline.text_encoder = mock_text_encoder # -- Mock transformer -- mock_transformer = MagicMock() mock_transformer.dtype = torch.bfloat16 mock_transformer.parameters = MagicMock(return_value=iter([torch.zeros(1)])) # encode_audio returns dict with audio_emb mock_transformer.encode_audio = MagicMock(return_value={"audio_emb": torch.zeros(1, 10, 64)}) # Forward returns noise prediction mock_transformer.return_value = (torch.zeros(1, 16, 20, 88, 128),) pipeline.transformer = mock_transformer # -- Mock VAE -- mock_vae = MagicMock() mock_vae.dtype = torch.bfloat16 mock_vae.config = MagicMock() mock_vae.config.scale_factor_temporal = 4 mock_vae.config.scale_factor_spatial = 8 mock_vae.config.latents_mean = [0.0] * 16 mock_vae.config.latents_std = [1.0] * 16 mock_vae.config.z_dim = 16 # encode returns mock with latent_dist mock_encode_result = MagicMock() mock_encode_result.latent_dist = MagicMock() mock_encode_result.latent_dist.mode = MagicMock(return_value=torch.zeros(1, 16, 1, 88, 128)) mock_vae.encode = MagicMock(return_value=mock_encode_result) # decode returns video mock_vae.decode = MagicMock(return_value=(torch.zeros(1, 3, 80, 704, 1024),)) pipeline.vae = mock_vae # -- Mock audio model -- mock_audio_model = MagicMock() mock_audio_model.device = torch.device("cpu") mock_audio_param = torch.zeros(1) mock_audio_model.parameters = MagicMock(return_value=iter([mock_audio_param])) mock_audio_model.return_value = MagicMock(hidden_states=[torch.zeros(1, 100, 1024)] * 25) pipeline.audio_model = mock_audio_model pipeline.audio_processor = MagicMock() pipeline.audio_processor.return_value = MagicMock(input_values=torch.zeros(1, 16000)) # -- Mock scheduler -- mock_scheduler = MagicMock() mock_scheduler.timesteps = torch.linspace(999, 0, 5) mock_scheduler.step = MagicMock(return_value=(torch.zeros(1, 16, 20, 88, 128),)) pipeline.scheduler = mock_scheduler pipeline._flow_shift = 3.0 # -- Bind methods from the real class -- pipeline.encode_prompt = Wan22S2VPipeline.encode_prompt.__get__(pipeline) pipeline.encode_ref_image = Wan22S2VPipeline.encode_ref_image.__get__(pipeline) pipeline.prepare_motion_latents = Wan22S2VPipeline.prepare_motion_latents.__get__(pipeline) pipeline.prepare_latents = Wan22S2VPipeline.prepare_latents.__get__(pipeline) pipeline.check_inputs = Wan22S2VPipeline.check_inputs.__get__(pipeline) pipeline.diffuse = Wan22S2VPipeline.diffuse.__get__(pipeline) pipeline._normalize_latents = Wan22S2VPipeline._normalize_latents.__get__(pipeline) pipeline._denormalize_latents = Wan22S2VPipeline._denormalize_latents.__get__(pipeline) pipeline._prompt_clean = Wan22S2VPipeline._prompt_clean # -- Build request -- sampling_params = _make_s2v_sampling( height=704, width=1024, num_frames=80, num_inference_steps=5, guidance_scale=4.5, generator=torch.Generator(device="cpu").manual_seed(42), max_sequence_length=512, ) ref_image = PIL.Image.new("RGB", (1024, 704)) audio_data = np.zeros(16000, dtype=np.float32) req = DiffusionRequestBatch( requests=[ SimpleNamespace( request_id="test", prompt={ "prompt": "test prompt", "negative_prompt": "bad quality", "multi_modal_data": {"image": ref_image, "audio": audio_data}, "additional_information": { "audio_path": audio_data, "pose_video": None, "init_first_frame": False, }, }, sampling_params=sampling_params, ) ] ) # -- Mock methods that use complex internal state -- pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 25, 64, 80), 1, 80)) # Mock progress bar context pipeline.progress_bar = MagicMock() pipeline.progress_bar.return_value.__enter__ = MagicMock(return_value=MagicMock()) pipeline.progress_bar.return_value.__exit__ = MagicMock(return_value=False) # Mock predict_noise_maybe_with_cfg from CFGParallelMixin pipeline.predict_noise_maybe_with_cfg = MagicMock(return_value=torch.zeros(1, 16, 20, 88, 128)) pipeline.predict_noise = Wan22S2VPipeline.predict_noise.__get__(pipeline) # Mock platform methods with patch("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.current_omni_platform") as mock_platform: mock_platform.empty_cache = MagicMock() mock_platform.is_available = MagicMock(return_value=False) with patch( "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.load_audio", return_value=(audio_data, 16000) ): result = Wan22S2VPipeline.forward(pipeline, req=req)[0] # -- Assertions -- # Text encoder was called mock_text_encoder.assert_called() # Audio encoding was invoked pipeline.encode_audio.assert_called_once() # VAE encode was called (for ref image and motion latents) mock_vae.encode.assert_called() # VAE decode was called mock_vae.decode.assert_called() # transformer.to("cpu") must NOT be called under HSDP mock_transformer.to.assert_not_called() # Output is a DiffusionOutput tuple with (video, audio_waveform, sample_rate) from vllm_omni.diffusion.data import DiffusionOutput assert isinstance(result, DiffusionOutput) video, audio_waveform, audio_sr = result.output assert video.shape[0] == 1 # batch 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.quantize_chunk def capture(chunk, value_range): transfers.append(chunk.shape[2]) return quantize(chunk, value_range) 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 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)