[Perf] Support request-level batching for Wan2.2 pipelines (#5676)

Signed-off-by: nagisa-kun <1434936049@qq.com>
Signed-off-by: nagisa <1434936049@qq.com>

tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py

@@ -1,14 +1,307 @@
+from types import SimpleNamespace
 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.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]
 
 
+def _make_s2v_sampling(**overrides):
+    values = {
+        "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)
+
+
+def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> 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.scheduler = MagicMock(timesteps=torch.tensor([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_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_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)
+
+    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)
@@ -260,31 +553,37 @@ def test_s2v_pipeline_hsdp_forward_complete_process():
     pipeline._prompt_clean = Wan22S2VPipeline._prompt_clean
 
     # -- Build request --
-    from vllm_omni.diffusion.request import OmniDiffusionRequest
-
-    sampling_params = MagicMock()
-    sampling_params.height = 704
-    sampling_params.width = 1024
-    sampling_params.num_frames = 80
-    sampling_params.num_inference_steps = 5
-    sampling_params.guidance_scale = 4.5
-    sampling_params.guidance_scale_provided = True
-    sampling_params.generator = None
-    sampling_params.seed = 42
+    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 = MagicMock(spec=OmniDiffusionRequest)
-    req.prompts = [
-        {
-            "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},
-        }
-    ]
-    req.sampling_params = sampling_params
+    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))
@@ -295,7 +594,7 @@ def test_s2v_pipeline_hsdp_forward_complete_process():
     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(16, 20, 88, 128))
+    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
@@ -306,7 +605,7 @@ def test_s2v_pipeline_hsdp_forward_complete_process():
         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)
+            result = Wan22S2VPipeline.forward(pipeline, req=req)[0]
 
     # -- Assertions --
     # Text encoder was called