[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_i2v_pipeline.py

@@ -8,15 +8,28 @@ import torch
 from PIL import Image
 from torch import nn
 
-from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar
+from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import (
     Wan22I2VPipeline,
+    get_wan22_i2v_post_process_func,
     get_wan22_i2v_pre_process_func,
 )
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
+def test_wan22_i2v_postprocess_honors_request_output_type() -> None:
+    video = torch.zeros(1, 4, 1, 2, 2)
+
+    output = get_wan22_i2v_post_process_func(SimpleNamespace())(
+        video,
+        sampling_params=SimpleNamespace(output_type="latent"),
+    )
+
+    assert output is video
+
+
 def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline:
     pipeline = object.__new__(Wan22I2VPipeline)
     nn.Module.__init__(pipeline)
@@ -31,6 +44,29 @@ def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline:
     return pipeline
 
 
+def _make_i2v_sampling(**overrides):
+    values = {
+        "height": 16,
+        "width": 16,
+        "num_frames": 5,
+        "num_inference_steps": 1,
+        "guidance_scale_provided": True,
+        "guidance_scale": 1.0,
+        "guidance_scale_2": None,
+        "guidance_scale_2_provided": False,
+        "boundary_ratio": None,
+        "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 test_i2v_preprocess_requires_image_and_resizes_to_480p_aspect() -> None:
     preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
     request = SimpleNamespace(
@@ -53,6 +89,88 @@ def test_i2v_preprocess_requires_image_and_resizes_to_480p_aspect() -> None:
         preprocess(missing_image)
 
 
+def _make_i2v_preprocess_request(last_image):
+    return SimpleNamespace(
+        prompt={
+            "prompt": "p",
+            "multi_modal_data": {
+                "image": Image.new("RGB", (320, 160), "red"),
+                "last_image": last_image,
+            },
+        },
+        sampling_params=SimpleNamespace(height=16, width=16),
+    )
+
+
+def test_i2v_preprocess_treats_empty_last_image_list_as_absent() -> None:
+    preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
+
+    result = preprocess(_make_i2v_preprocess_request([]))
+
+    assert result.prompt["multi_modal_data"]["last_image"] is None
+    assert result.batch_compatibility_key == ("wan22_i2v_last_image", False)
+
+
+def test_i2v_preprocess_unwraps_single_last_image_list() -> None:
+    preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
+    last_image = Image.new("RGB", (16, 16), "blue")
+
+    result = preprocess(_make_i2v_preprocess_request([last_image]))
+
+    assert result.prompt["multi_modal_data"]["last_image"] is last_image
+    assert result.batch_compatibility_key == ("wan22_i2v_last_image", True)
+
+
+@pytest.mark.parametrize(
+    "last_image",
+    [
+        Image.new("RGB", (16, 16), "blue"),
+        torch.zeros(3, 16, 16),
+    ],
+)
+def test_i2v_preprocess_preserves_supported_last_image(last_image) -> None:
+    preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
+
+    result = preprocess(_make_i2v_preprocess_request(last_image))
+
+    assert result.prompt["multi_modal_data"]["last_image"] is last_image
+    assert result.batch_compatibility_key == ("wan22_i2v_last_image", True)
+
+
+def test_i2v_preprocess_loads_last_image_path(tmp_path) -> None:
+    preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
+    path = tmp_path / "last.png"
+    Image.new("RGB", (16, 16), "blue").save(path)
+
+    result = preprocess(_make_i2v_preprocess_request(str(path)))
+
+    last_image = result.prompt["multi_modal_data"]["last_image"]
+    assert isinstance(last_image, Image.Image)
+    assert last_image.mode == "RGB"
+    assert result.batch_compatibility_key == ("wan22_i2v_last_image", True)
+
+
+def test_i2v_preprocess_rejects_multiple_last_images() -> None:
+    preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
+
+    with pytest.raises(ValueError, match="at most one last_image"):
+        preprocess(
+            _make_i2v_preprocess_request(
+                [
+                    Image.new("RGB", (16, 16), "blue"),
+                    Image.new("RGB", (16, 16), "green"),
+                ]
+            )
+        )
+
+
+def test_i2v_preprocess_rejects_unsupported_last_image_type() -> None:
+    preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
+
+    with pytest.raises(TypeError, match="Unsupported last_image format"):
+        preprocess(_make_i2v_preprocess_request({"not": "an image"}))
+
+
 def test_i2v_diffuse_selects_stage_guidance_and_expands_timesteps() -> None:
     pipeline = _make_i2v_pipeline(expand_timesteps=True)
     latents = torch.zeros(1, 4, 2, 4, 4)
@@ -126,3 +244,169 @@ def test_i2v_prepare_latents_builds_expand_condition_and_first_frame_mask() -> N
     assert first_frame_mask.shape == (1, 1, 2, 2, 2)
     assert first_frame_mask[:, :, 0].sum() == 0
     assert first_frame_mask[:, :, 1].sum() == 4
+
+
+def test_i2v_prepare_latents_preserves_batched_image_conditions() -> None:
+    pipeline = _make_i2v_pipeline(expand_timesteps=True)
+    generators = [
+        torch.Generator(device="cpu").manual_seed(1),
+        torch.Generator(device="cpu").manual_seed(2),
+    ]
+
+    latents, condition, first_frame_mask = pipeline.prepare_latents(
+        image=torch.zeros(2, 3, 16, 16),
+        batch_size=2,
+        num_channels_latents=4,
+        height=16,
+        width=16,
+        num_frames=5,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+        generator=generators,
+    )
+
+    assert latents.shape == (2, 4, 2, 2, 2)
+    assert condition.shape == (2, 4, 1, 2, 2)
+    assert first_frame_mask.shape == (2, 1, 2, 2, 2)
+
+
+def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch) -> None:
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v.current_omni_platform",
+        SimpleNamespace(is_available=lambda: False),
+    )
+    pipeline = _make_i2v_pipeline(expand_timesteps=True)
+    pipeline.scheduler = StubScheduler([9])
+    pipeline.od_config = SimpleNamespace(flow_shift=5.0)
+    pipeline._sample_solver = "unipc"
+    pipeline._flow_shift = 5.0
+    pipeline.boundary_ratio = 0.875
+    pipeline.has_image_encoder = False
+    pipeline._guidance_scale = None
+    pipeline._guidance_scale_2 = None
+    pipeline._num_timesteps = None
+    pipeline._current_timestep = None
+    pipeline.check_inputs = lambda **kwargs: None
+    prepare_call = {}
+
+    def _fake_encode_prompt(**kwargs):
+        batch_size = len(kwargs["prompt"])
+        n = batch_size * kwargs["num_videos_per_prompt"]
+        return torch.zeros(n, 2, 3), None
+
+    def _fake_prepare_latents(**kwargs):
+        prepare_call.update(kwargs)
+        batch_size = kwargs["batch_size"]
+        return (
+            kwargs["latents"],
+            torch.zeros(batch_size, 4, 1, 2, 2),
+            torch.ones(batch_size, 1, 2, 2, 2),
+        )
+
+    pipeline.encode_prompt = _fake_encode_prompt  # type: ignore[method-assign]
+    pipeline.prepare_latents = _fake_prepare_latents  # type: ignore[method-assign]
+    pipeline.diffuse = lambda **kwargs: kwargs["latents"]  # type: ignore[method-assign]
+
+    gen_a = torch.Generator(device="cpu").manual_seed(1)
+    gen_b = torch.Generator(device="cpu").manual_seed(2)
+    latents_a = torch.zeros(2, 4, 2, 2, 2)
+    latents_b = torch.ones(2, 4, 2, 2, 2)
+    image_a = torch.zeros(1, 3, 16, 16)
+    image_b = torch.ones(1, 3, 16, 16)
+    batch = DiffusionRequestBatch(
+        requests=[
+            SimpleNamespace(
+                request_id="a",
+                prompt={"prompt": "first", "multi_modal_data": {"image": image_a}},
+                sampling_params=_make_i2v_sampling(
+                    generator=gen_a,
+                    latents=latents_a,
+                    num_outputs_per_prompt=2,
+                ),
+            ),
+            SimpleNamespace(
+                request_id="b",
+                prompt={"prompt": "second", "multi_modal_data": {"image": image_b}},
+                sampling_params=_make_i2v_sampling(
+                    generator=gen_b,
+                    latents=latents_b,
+                    num_outputs_per_prompt=2,
+                ),
+            ),
+        ]
+    )
+
+    outputs = pipeline.forward(batch)
+
+    assert prepare_call["batch_size"] == 4
+    assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b]
+    torch.testing.assert_close(prepare_call["latents"], torch.cat([latents_a, latents_b]))
+    torch.testing.assert_close(
+        prepare_call["image"],
+        torch.cat([image_a, image_a, image_b, image_b]),
+    )
+    assert len(outputs) == 2
+    torch.testing.assert_close(outputs[0].output, latents_a)
+    torch.testing.assert_close(outputs[1].output, latents_b)
+
+
+def test_i2v_forward_rejects_mismatched_tensor_condition_shapes(monkeypatch) -> None:
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v.current_omni_platform",
+        SimpleNamespace(is_available=lambda: False),
+    )
+    pipeline = _make_i2v_pipeline(expand_timesteps=True)
+    pipeline.scheduler = StubScheduler([9])
+    pipeline.od_config = SimpleNamespace(flow_shift=5.0)
+    pipeline._sample_solver = "unipc"
+    pipeline._flow_shift = 5.0
+    pipeline.boundary_ratio = 0.875
+    pipeline.has_image_encoder = False
+    pipeline._guidance_scale = None
+    pipeline._guidance_scale_2 = None
+    pipeline._num_timesteps = None
+    pipeline._current_timestep = None
+    pipeline.check_inputs = lambda **kwargs: None
+    pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None)  # type: ignore[method-assign]
+    batch = DiffusionRequestBatch(
+        requests=[
+            SimpleNamespace(
+                request_id="a",
+                prompt={"prompt": "first", "multi_modal_data": {"image": torch.zeros(1, 3, 16, 16)}},
+                sampling_params=_make_i2v_sampling(),
+            ),
+            SimpleNamespace(
+                request_id="b",
+                prompt={"prompt": "second", "multi_modal_data": {"image": torch.zeros(1, 3, 16, 8)}},
+                sampling_params=_make_i2v_sampling(),
+            ),
+        ]
+    )
+
+    with pytest.raises(ValueError, match="image condition"):
+        pipeline.forward(batch)
+
+
+def test_i2v_forward_rejects_mixed_last_image_presence() -> None:
+    pipeline = _make_i2v_pipeline(expand_timesteps=True)
+    image = torch.zeros(1, 3, 16, 16)
+    batch = DiffusionRequestBatch(
+        requests=[
+            SimpleNamespace(
+                request_id="a",
+                prompt={
+                    "prompt": "first",
+                    "multi_modal_data": {"image": image, "last_image": image},
+                },
+                sampling_params=_make_i2v_sampling(),
+            ),
+            SimpleNamespace(
+                request_id="b",
+                prompt={"prompt": "second", "multi_modal_data": {"image": image}},
+                sampling_params=_make_i2v_sampling(),
+            ),
+        ]
+    )
+
+    with pytest.raises(ValueError, match="mix of provided and missing last_image"):
+        pipeline.forward(batch)

tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py

@@ -93,6 +93,29 @@ def _make_pipeline() -> Wan22Pipeline:
     return pipeline
 
 
+def _make_sampling(**overrides):
+    values = {
+        "height": None,
+        "width": None,
+        "num_frames": 1,
+        "num_inference_steps": 2,
+        "guidance_scale_provided": True,
+        "guidance_scale": 1.0,
+        "guidance_scale_2": None,
+        "guidance_scale_2_provided": False,
+        "boundary_ratio": None,
+        "generator": None,
+        "seed": None,
+        "num_outputs_per_prompt": 1,
+        "max_sequence_length": 32,
+        "latents": None,
+        "output_type": "latent",
+        "extra_args": {},
+    }
+    values.update(overrides)
+    return SimpleNamespace(**values)
+
+
 @pytest.mark.parametrize(
     ("sampling_params_kwargs", "expected_low", "expected_high"),
     [
@@ -129,9 +152,10 @@ def test_forward_delegates_denoising_to_diffuse(
     )
     batch = DiffusionRequestBatch(requests=[mock_req])
 
-    output = pipeline.forward(batch)
+    outputs = pipeline.forward(batch)
 
-    assert torch.equal(output.output, torch.ones((1, 4, 1, 8, 8)))
+    assert len(outputs) == 1
+    assert torch.equal(outputs[0].output, torch.ones((1, 4, 1, 8, 8)))
     assert torch.equal(captured["prompt_embeds"], torch.zeros(1, 32, 8))
     assert torch.equal(captured["timesteps"], pipeline.scheduler.timesteps)
     assert captured["guidance_low"] == expected_low
@@ -142,6 +166,134 @@ def test_forward_delegates_denoising_to_diffuse(
     assert pipeline.scheduler.set_timesteps_calls == [(2, torch.device("cpu"))]
 
 
+def test_forward_batches_text_generators_latents_and_splits_outputs() -> None:
+    pipeline = _make_pipeline()
+    encode_call = {}
+    prepare_call = {}
+
+    def _fake_encode_prompt(**kwargs):
+        encode_call.update(kwargs)
+        batch_size = len(kwargs["prompt"])
+        n = batch_size * kwargs["num_videos_per_prompt"]
+        return torch.arange(n, dtype=torch.float32).view(n, 1, 1), torch.zeros(n, 1, 1)
+
+    def _fake_prepare_latents(**kwargs):
+        prepare_call.update(kwargs)
+        return kwargs["latents"]
+
+    pipeline.encode_prompt = _fake_encode_prompt  # type: ignore[method-assign]
+    pipeline.prepare_latents = _fake_prepare_latents  # type: ignore[method-assign]
+    pipeline.diffuse = lambda **kwargs: kwargs["latents"]  # type: ignore[method-assign]
+
+    gen_a = torch.Generator(device="cpu").manual_seed(1)
+    gen_b = torch.Generator(device="cpu").manual_seed(2)
+    latents_a = torch.zeros(2, 4, 1, 2, 2)
+    latents_b = torch.ones(2, 4, 1, 2, 2)
+    batch = DiffusionRequestBatch(
+        requests=[
+            SimpleNamespace(
+                request_id="a",
+                prompt={"prompt": "first", "negative_prompt": "bad first"},
+                sampling_params=_make_sampling(
+                    generator=gen_a,
+                    latents=latents_a,
+                    num_outputs_per_prompt=2,
+                ),
+            ),
+            SimpleNamespace(
+                request_id="b",
+                prompt={"prompt": "second", "negative_prompt": "bad second"},
+                sampling_params=_make_sampling(
+                    generator=gen_b,
+                    latents=latents_b,
+                    num_outputs_per_prompt=2,
+                ),
+            ),
+        ]
+    )
+
+    outputs = pipeline.forward(batch)
+
+    assert encode_call["prompt"] == ["first", "second"]
+    assert encode_call["negative_prompt"] == ["bad first", "bad second"]
+    assert prepare_call["batch_size"] == 4
+    assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b]
+    torch.testing.assert_close(prepare_call["latents"], torch.cat([latents_a, latents_b]))
+    assert len(outputs) == 2
+    torch.testing.assert_close(outputs[0].output, latents_a)
+    torch.testing.assert_close(outputs[1].output, latents_b)
+
+
+def test_forward_batches_precomputed_prompt_embeddings() -> None:
+    pipeline = _make_pipeline()
+    diffuse_call = {}
+    pipeline.encode_prompt = lambda **kwargs: pytest.fail("text encoder must not run")  # type: ignore[method-assign]
+    pipeline.prepare_latents = lambda **kwargs: torch.zeros(kwargs["batch_size"], 4, 1, 2, 2)  # type: ignore[method-assign]
+
+    def _fake_diffuse(**kwargs):
+        diffuse_call.update(kwargs)
+        return kwargs["latents"]
+
+    pipeline.diffuse = _fake_diffuse  # type: ignore[method-assign]
+    embeds_a = torch.zeros(3, 4)
+    embeds_b = torch.ones(3, 4)
+    negative_a = torch.full((3, 4), 2.0)
+    negative_b = torch.full((3, 4), 3.0)
+    batch = DiffusionRequestBatch(
+        requests=[
+            SimpleNamespace(
+                request_id="a",
+                prompt={"prompt_embeds": embeds_a, "negative_prompt_embeds": negative_a},
+                sampling_params=_make_sampling(guidance_scale=4.0),
+            ),
+            SimpleNamespace(
+                request_id="b",
+                prompt={"prompt_embeds": embeds_b, "negative_prompt_embeds": negative_b},
+                sampling_params=_make_sampling(guidance_scale=4.0),
+            ),
+        ]
+    )
+
+    outputs = pipeline.forward(batch)
+
+    torch.testing.assert_close(diffuse_call["prompt_embeds"], torch.stack([embeds_a, embeds_b]))
+    torch.testing.assert_close(diffuse_call["negative_prompt_embeds"], torch.stack([negative_a, negative_b]))
+    assert len(outputs) == 2
+
+
+def test_prepare_latents_with_request_generators_matches_single_generation() -> None:
+    pipeline = _make_pipeline()
+    kwargs = {
+        "num_channels_latents": 4,
+        "height": 16,
+        "width": 16,
+        "num_frames": 5,
+        "dtype": torch.float32,
+        "device": torch.device("cpu"),
+    }
+
+    batched = Wan22Pipeline.prepare_latents(
+        pipeline,
+        batch_size=2,
+        generator=[torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)],
+        **kwargs,
+    )
+    singles = torch.cat(
+        [
+            Wan22Pipeline.prepare_latents(
+                pipeline,
+                batch_size=1,
+                generator=torch.Generator().manual_seed(seed),
+                **kwargs,
+            )
+            for seed in (1, 2)
+        ]
+    )
+
+    torch.testing.assert_close(batched, singles)
+    assert not torch.equal(batched[0], batched[1])
+
+
 def test_diffuse_runs_prediction_and_scheduler_for_each_timestep() -> None:
     pipeline = _make_pipeline()
     latents = torch.zeros((1, 1, 1, 2, 2), dtype=torch.float32)

tests/diffusion/models/wan2_2/test_wan22_pipeline_helpers.py

@@ -10,6 +10,7 @@ import torch
 import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 as wan22_module
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
     create_transformer_from_config,
+    get_wan22_post_process_func,
     load_transformer_config,
     retrieve_latents,
 )
@@ -26,6 +27,17 @@ class _LatentDist:
         return torch.tensor([2.0])
 
 
+def test_wan22_postprocess_honors_request_output_type() -> None:
+    video = torch.zeros(1, 4, 1, 2, 2)
+
+    output = get_wan22_post_process_func(SimpleNamespace())(
+        video,
+        sampling_params=SimpleNamespace(output_type="latent"),
+    )
+
+    assert output is video
+
+
 def test_retrieve_latents_supports_sample_mode_argmax_and_direct_latents() -> None:
     generator = torch.Generator(device="cpu")
 

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 = _make_s2v_sampling(
-
+        height=704,
-    sampling_params = MagicMock()
+        width=1024,
-    sampling_params.height = 704
+        num_frames=80,
-    sampling_params.width = 1024
+        num_inference_steps=5,
-    sampling_params.num_frames = 80
+        guidance_scale=4.5,
-    sampling_params.num_inference_steps = 5
+        generator=torch.Generator(device="cpu").manual_seed(42),
-    sampling_params.guidance_scale = 4.5
+        max_sequence_length=512,
-    sampling_params.guidance_scale_provided = True
+    )
-    sampling_params.generator = None
-    sampling_params.seed = 42
 
     ref_image = PIL.Image.new("RGB", (1024, 704))
     audio_data = np.zeros(16000, dtype=np.float32)
 
-    req = MagicMock(spec=OmniDiffusionRequest)
+    req = DiffusionRequestBatch(
-    req.prompts = [
+        requests=[
-        {
+            SimpleNamespace(
-            "prompt": "test prompt",
+                request_id="test",
-            "negative_prompt": "bad quality",
+                prompt={
-            "multi_modal_data": {"image": ref_image, "audio": audio_data},
+                    "prompt": "test prompt",
-            "additional_information": {"audio_path": audio_data, "pose_video": None, "init_first_frame": False},
+                    "negative_prompt": "bad quality",
-        }
+                    "multi_modal_data": {"image": ref_image, "audio": audio_data},
-    ]
+                    "additional_information": {
-    req.sampling_params = sampling_params
+                        "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

tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py

@@ -8,13 +8,14 @@ import torch
 from PIL import Image
 from torch import nn
 
-from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar
+from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace import (
     Wan22VACEPipeline,
     create_vace_transformer_from_config,
     get_wan22_vace_pre_process_func,
 )
 from vllm_omni.diffusion.models.wan2_2.wan2_2_vace_transformer import WanVACETransformer3DModel
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
@@ -32,6 +33,27 @@ def _make_vace_pipeline() -> Wan22VACEPipeline:
     return pipeline
 
 
+def _make_vace_sampling(**overrides):
+    values = {
+        "height": 16,
+        "width": 16,
+        "num_frames": 5,
+        "num_inference_steps": 1,
+        "guidance_scale_provided": True,
+        "guidance_scale": 1.0,
+        "boundary_ratio": None,
+        "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 test_vace_preprocess_collects_reference_video_and_mask_inputs() -> None:
     preprocess = get_wan22_vace_pre_process_func(SimpleNamespace())
     ref = Image.new("RGB", (320, 160), "green")
@@ -138,3 +160,189 @@ def test_vace_diffuse_passes_context_and_scale_to_cfg_branches() -> None:
     assert calls[0]["positive_kwargs"]["vace_context"] is vace_context
     assert calls[0]["negative_kwargs"]["vace_context_scale"] == 0.75
     torch.testing.assert_close(result, torch.ones_like(latents))
+
+
+def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> None:
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform",
+        SimpleNamespace(is_available=lambda: False),
+    )
+    pipeline = _make_vace_pipeline()
+    pipeline.transformer.vace_patch_embedding = None
+    pipeline.transformer_2 = None
+    pipeline.scheduler = StubScheduler([9])
+    pipeline.od_config = SimpleNamespace(flow_shift=3.0)
+    pipeline._sample_solver = "unipc"
+    pipeline._flow_shift = 3.0
+    pipeline.boundary_ratio = None
+    pipeline._guidance_scale = None
+    pipeline._num_timesteps = None
+    pipeline._current_timestep = None
+    pipeline.check_inputs = lambda **kwargs: None
+    prepare_call = {}
+
+    def _fake_encode_prompt(**kwargs):
+        batch_size = len(kwargs["prompt"])
+        n = batch_size * kwargs["num_videos_per_prompt"]
+        return torch.zeros(n, 2, 3), None
+
+    def _fake_prepare_latents(**kwargs):
+        prepare_call.update(kwargs)
+        return kwargs["latents"]
+
+    pipeline.encode_prompt = _fake_encode_prompt  # type: ignore[method-assign]
+    pipeline.prepare_latents = _fake_prepare_latents  # type: ignore[method-assign]
+    pipeline.diffuse = lambda **kwargs: kwargs["latents"]  # type: ignore[method-assign]
+    gen_a = torch.Generator(device="cpu").manual_seed(1)
+    gen_b = torch.Generator(device="cpu").manual_seed(2)
+    latents_a = torch.zeros(2, 4, 2, 2, 2)
+    latents_b = torch.ones(2, 4, 2, 2, 2)
+    batch = DiffusionRequestBatch(
+        requests=[
+            SimpleNamespace(
+                request_id="a",
+                prompt={"prompt": "first"},
+                sampling_params=_make_vace_sampling(
+                    generator=gen_a,
+                    latents=latents_a,
+                    num_outputs_per_prompt=2,
+                ),
+            ),
+            SimpleNamespace(
+                request_id="b",
+                prompt={"prompt": "second"},
+                sampling_params=_make_vace_sampling(
+                    generator=gen_b,
+                    latents=latents_b,
+                    num_outputs_per_prompt=2,
+                ),
+            ),
+        ]
+    )
+
+    outputs = pipeline.forward(batch)
+
+    assert prepare_call["batch_size"] == 4
+    assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b]
+    torch.testing.assert_close(prepare_call["latents"], torch.cat([latents_a, latents_b]))
+    assert len(outputs) == 2
+    torch.testing.assert_close(outputs[0].output, latents_a)
+    torch.testing.assert_close(outputs[1].output, latents_b)
+
+
+def test_vace_forward_rejects_mismatched_condition_shapes(monkeypatch) -> None:
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform",
+        SimpleNamespace(is_available=lambda: False),
+    )
+    pipeline = _make_vace_pipeline()
+    pipeline.transformer.vace_patch_embedding = object()
+    pipeline.transformer_2 = None
+    pipeline.scheduler = StubScheduler([9])
+    pipeline.od_config = SimpleNamespace(flow_shift=3.0)
+    pipeline._sample_solver = "unipc"
+    pipeline._flow_shift = 3.0
+    pipeline.boundary_ratio = None
+    pipeline._guidance_scale = None
+    pipeline._num_timesteps = None
+    pipeline._current_timestep = None
+    pipeline.check_inputs = lambda **kwargs: None
+    pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None)  # type: ignore[method-assign]
+
+    def _fake_preprocess(video, **kwargs):
+        del kwargs
+        mask = torch.ones_like(video)
+        return video, mask, [[]]
+
+    pipeline.preprocess_conditions = _fake_preprocess  # type: ignore[method-assign]
+    pipeline.prepare_video_latents = (  # type: ignore[method-assign]
+        lambda video, mask, refs, generator, device: torch.zeros(1, 8, video.shape[2], 2, 2)
+    )
+    pipeline.prepare_masks = lambda mask, refs: torch.zeros(1, 4, mask.shape[2], 2, 2)  # type: ignore[method-assign]
+    batch = DiffusionRequestBatch(
+        requests=[
+            SimpleNamespace(
+                request_id="a",
+                prompt={
+                    "prompt": "first",
+                    "additional_information": {"source_video": torch.zeros(1, 3, 5, 16, 16)},
+                },
+                sampling_params=_make_vace_sampling(),
+            ),
+            SimpleNamespace(
+                request_id="b",
+                prompt={
+                    "prompt": "second",
+                    "additional_information": {"source_video": torch.zeros(1, 3, 3, 16, 16)},
+                },
+                sampling_params=_make_vace_sampling(),
+            ),
+        ]
+    )
+
+    with pytest.raises(ValueError, match="VACE condition"):
+        pipeline.forward(batch)
+
+
+def test_vace_forward_rejects_mismatched_reference_image_counts(monkeypatch) -> None:
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform",
+        SimpleNamespace(is_available=lambda: False),
+    )
+    pipeline = _make_vace_pipeline()
+    pipeline.transformer.vace_patch_embedding = object()
+    pipeline.transformer_2 = None
+    pipeline.scheduler = StubScheduler([9])
+    pipeline.od_config = SimpleNamespace(flow_shift=3.0)
+    pipeline._sample_solver = "unipc"
+    pipeline._flow_shift = 3.0
+    pipeline.boundary_ratio = None
+    pipeline._guidance_scale = None
+    pipeline._num_timesteps = None
+    pipeline._current_timestep = None
+    pipeline.check_inputs = lambda **kwargs: None
+    pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None)  # type: ignore[method-assign]
+
+    def _fake_preprocess(video, mask, reference_images, **kwargs):
+        del kwargs
+        return video, mask, [reference_images]
+
+    pipeline.preprocess_conditions = _fake_preprocess  # type: ignore[method-assign]
+    pipeline.prepare_video_latents = (  # type: ignore[method-assign]
+        lambda video, mask, refs, generator, device: torch.zeros(1, 8, 5, 2, 2)
+    )
+    pipeline.prepare_masks = lambda mask, refs: torch.zeros(1, 4, 5, 2, 2)  # type: ignore[method-assign]
+    video = torch.zeros(1, 3, 5, 16, 16)
+    mask = torch.ones_like(video)
+    reference = torch.zeros(3, 16, 16)
+    batch = DiffusionRequestBatch(
+        requests=[
+            SimpleNamespace(
+                request_id="a",
+                prompt={
+                    "prompt": "first",
+                    "additional_information": {
+                        "source_video": video,
+                        "mask": mask,
+                        "reference_images": [reference],
+                    },
+                },
+                sampling_params=_make_vace_sampling(),
+            ),
+            SimpleNamespace(
+                request_id="b",
+                prompt={
+                    "prompt": "second",
+                    "additional_information": {
+                        "source_video": video,
+                        "mask": mask,
+                        "reference_images": [reference, reference],
+                    },
+                },
+                sampling_params=_make_vace_sampling(),
+            ),
+        ]
+    )
+
+    with pytest.raises(ValueError, match="same number of reference images"):
+        pipeline.forward(batch)

tests/diffusion/test_diffusion_engine_rpc_routing.py

@@ -123,6 +123,7 @@ class _ConcurrencyTrackingExecutor:
 def _make_request(tag: str):
     sampling_params = dict(_SAMPLING_KEY_DEFAULTS)
     sampling_params["num_inference_steps"] = 1
+    sampling_params["extra_args"] = {}
     return SimpleNamespace(
         request_id=tag,
         prompt=f"prompt_{tag}",

tests/diffusion/test_diffusion_scheduler.py

@@ -217,13 +217,21 @@ class TestGetRequestBatchSamplingParamsKey:
         num_inference_steps: int = 2,
         seed: int | None = 123,
         generator: torch.Generator | None = None,
+        extra_args: dict | None = None,
+        condition_key: tuple | None = None,
     ) -> OmniDiffusionRequest:
         sp = OmniDiffusionSamplingParams(
             num_inference_steps=num_inference_steps,
             seed=seed,
             generator=generator,
+            extra_args=extra_args or {},
+        )
+        return OmniDiffusionRequest(
+            prompt="prompt",
+            sampling_params=sp,
+            request_id=f"req-{num_inference_steps}",
+            batch_compatibility_key=condition_key,
         )
-        return OmniDiffusionRequest(prompt="prompt", sampling_params=sp, request_id=f"req-{num_inference_steps}")
 
     def test_distinguishes_num_inference_steps(self) -> None:
         scheduler = RequestScheduler()
@@ -240,6 +248,57 @@ class TestGetRequestBatchSamplingParamsKey:
             self._make(seed=1, generator=gen_a)
         ) == scheduler._build_sampling_params_key(self._make(seed=2, generator=gen_b))
 
+    @pytest.mark.parametrize(
+        ("first_extra_args", "second_extra_args"),
+        [
+            ({"sample_solver": "unipc"}, {"sample_solver": "euler"}),
+            ({"flow_shift": 3.0}, {"flow_shift": 5.0}),
+        ],
+    )
+    def test_distinguishes_wan_scheduler_structure(
+        self,
+        first_extra_args: dict,
+        second_extra_args: dict,
+    ) -> None:
+        scheduler = RequestScheduler()
+
+        assert scheduler._build_sampling_params_key(
+            self._make(extra_args=first_extra_args)
+        ) != scheduler._build_sampling_params_key(self._make(extra_args=second_extra_args))
+
+    @pytest.mark.parametrize(
+        ("first_extra_args", "second_extra_args"),
+        [
+            ({"sample_solver": " Euler "}, {"sample_solver": "euler"}),
+            ({"flow_shift": "5.0"}, {"flow_shift": 5.0}),
+        ],
+    )
+    def test_normalizes_equivalent_wan_scheduler_structure(
+        self,
+        first_extra_args: dict,
+        second_extra_args: dict,
+    ) -> None:
+        scheduler = RequestScheduler()
+
+        assert scheduler._build_sampling_params_key(
+            self._make(extra_args=first_extra_args)
+        ) == scheduler._build_sampling_params_key(self._make(extra_args=second_extra_args))
+
+    def test_uses_none_for_unspecified_wan_scheduler_structure(self) -> None:
+        scheduler = RequestScheduler()
+
+        key = scheduler._build_sampling_params_key(self._make())
+
+        assert key.sample_solver is None
+        assert key.flow_shift is None
+
+    def test_distinguishes_pipeline_condition_structure(self) -> None:
+        scheduler = RequestScheduler()
+
+        assert scheduler._build_sampling_params_key(
+            self._make(condition_key=("wan22_s2v_condition", True))
+        ) != scheduler._build_sampling_params_key(self._make(condition_key=("wan22_s2v_condition", False)))
+
 
 class TestRequestScheduler:
     def setup_method(self) -> None:
@@ -602,6 +661,38 @@ class TestRequestScheduler:
         assert first.num_running_reqs == 1
         assert first.num_waiting_reqs == 1
 
+    def test_batches_incompatible_pipeline_conditions_separately(self) -> None:
+        scheduler = RequestScheduler()
+        scheduler.initialize(SimpleNamespace(max_num_seqs=2))
+
+        request_a = OmniDiffusionRequest(
+            prompt="a",
+            sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, seed=123),
+            request_id="a",
+            batch_compatibility_key=("wan22_s2v_condition", True),
+        )
+        request_b = OmniDiffusionRequest(
+            prompt="b",
+            sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, seed=456),
+            request_id="b",
+            batch_compatibility_key=("wan22_s2v_condition", False),
+        )
+        scheduler.add_request(request_a)
+        scheduler.add_request(request_b)
+
+        first = scheduler.schedule()
+
+        assert _new_ids(first) == ["a"]
+        assert first.num_running_reqs == 1
+        assert first.num_waiting_reqs == 1
+
+        scheduler.update_from_output(first, _make_request_output("a"))
+        second = scheduler.schedule()
+
+        assert _new_ids(second) == ["b"]
+        assert second.num_running_reqs == 1
+        assert second.num_waiting_reqs == 0
+
     def test_batches_different_request_local_seed_together(self) -> None:
         scheduler = RequestScheduler()
         scheduler.initialize(SimpleNamespace(max_num_seqs=2))
@@ -625,6 +716,47 @@ class TestRequestScheduler:
         assert first.num_running_reqs == 2
         assert first.num_waiting_reqs == 0
 
+    @pytest.mark.parametrize(
+        ("first_extra_args", "second_extra_args"),
+        [
+            ({"sample_solver": "unipc"}, {"sample_solver": "euler"}),
+            ({"flow_shift": 3.0}, {"flow_shift": 5.0}),
+            ({"sample_solver": "unipc"}, {}),
+            ({"flow_shift": 3.0}, {}),
+        ],
+    )
+    def test_batches_incompatible_wan_scheduler_structure_separately(
+        self,
+        first_extra_args: dict,
+        second_extra_args: dict,
+    ) -> None:
+        scheduler = RequestScheduler()
+        scheduler.initialize(SimpleNamespace(max_num_seqs=2))
+        first_id = scheduler.add_request(
+            _make_step_request(
+                "a",
+                sampling_params=OmniDiffusionSamplingParams(
+                    num_inference_steps=2,
+                    extra_args=first_extra_args,
+                ),
+            )
+        )
+        scheduler.add_request(
+            _make_step_request(
+                "b",
+                sampling_params=OmniDiffusionSamplingParams(
+                    num_inference_steps=2,
+                    extra_args=second_extra_args,
+                ),
+            )
+        )
+
+        first = scheduler.schedule()
+
+        assert _new_ids(first) == [first_id]
+        assert first.num_running_reqs == 1
+        assert first.num_waiting_reqs == 1
+
     def test_incompatible_waiting_head_blocks_later_compatible_request(self) -> None:
         scheduler = RequestScheduler()
         scheduler.initialize(SimpleNamespace(max_num_seqs=3))

tests/engine/test_async_omni_engine_stage_init.py

@@ -399,6 +399,101 @@ def test_initialize_local_llm_replica_scopes_runtime_env(monkeypatch):
     assert os.environ[runtime_env_var] == "parent-value"
 
 
+def test_initialize_diffusion_stage_applies_client_batch_size_to_engine(monkeypatch):
+    import vllm_omni.diffusion.stage_diffusion_client as client_mod
+    import vllm_omni.engine.stage_init_utils as init_mod
+
+    od_config = types.SimpleNamespace(max_num_seqs=1)
+    captured: dict[str, object] = {}
+    monkeypatch.setattr(init_mod, "build_diffusion_config", lambda *args: od_config)
+
+    def _capture_client(model, config, metadata, stage_init_timeout, batch_size, use_inline):
+        captured.update(
+            model=model,
+            config=config,
+            metadata=metadata,
+            stage_init_timeout=stage_init_timeout,
+            batch_size=batch_size,
+            use_inline=use_inline,
+        )
+        return object()
+
+    monkeypatch.setattr(client_mod, "create_diffusion_client", _capture_client)
+    metadata = _make_diffusion_metadata(0)
+
+    init_mod.initialize_diffusion_stage(
+        0,
+        "dummy-model",
+        types.SimpleNamespace(),
+        metadata,
+        stage_init_timeout=12,
+        batch_size=4,
+        use_inline=True,
+    )
+
+    assert od_config.max_num_seqs == 4
+    assert captured == {
+        "model": "dummy-model",
+        "config": od_config,
+        "metadata": metadata,
+        "stage_init_timeout": 12,
+        "batch_size": 4,
+        "use_inline": True,
+    }
+
+
+def test_launch_diffusion_stage_replica_applies_batch_size_to_config(monkeypatch):
+    import vllm_omni.diffusion.stage_diffusion_client as client_mod
+    import vllm_omni.diffusion.stage_diffusion_proc as proc_mod
+    import vllm_omni.engine.stage_engine_startup as startup_mod
+
+    od_config = types.SimpleNamespace(max_num_seqs=1, parallel_config=types.SimpleNamespace(world_size=1))
+    monkeypatch.setattr(startup_mod, "build_diffusion_config", lambda *args: od_config)
+    monkeypatch.setattr(startup_mod, "acquire_device_locks", lambda *args: [])
+    monkeypatch.setattr(
+        startup_mod,
+        "register_stage_with_omni_master",
+        lambda **kwargs: types.SimpleNamespace(
+            handshake_address="tcp://127.0.0.1:26001",
+            input_address="tcp://127.0.0.1:26002",
+            output_address="tcp://127.0.0.1:26003",
+        ),
+    )
+
+    omni_master_server = types.SimpleNamespace(
+        address="127.0.0.1",
+        port=25000,
+        release_route_port_reservations=lambda *args, **kwargs: None,
+    )
+    proc_manager = types.SimpleNamespace(
+        addresses=types.SimpleNamespace(
+            inputs=["tcp://127.0.0.1:26002"],
+            outputs=["tcp://127.0.0.1:26003"],
+        )
+    )
+    monkeypatch.setattr(proc_mod, "StageDiffusionProcManager", lambda **kwargs: proc_manager)
+    sentinel_client = object()
+    monkeypatch.setattr(
+        client_mod.StageDiffusionClient,
+        "from_addresses",
+        lambda metadata, **kwargs: sentinel_client,
+    )
+
+    result, resources = startup_mod.launch_diffusion_stage_replica(
+        model="dummy-model",
+        stage_config=types.SimpleNamespace(),
+        metadata=types.SimpleNamespace(stage_id=0),
+        stage_init_timeout=12,
+        batch_size=4,
+        use_inline=False,
+        omni_master_server=omni_master_server,
+    )
+
+    assert result is sentinel_client
+    assert od_config.max_num_seqs == 4
+    assert resources.manager is proc_manager
+
+
 def test_stage_runtime_initializes_stage_pools(monkeypatch):
     import vllm_omni.engine.stage_runtime as runtime_mod
 

tests/engine/test_single_stage_mode.py

@@ -1138,7 +1138,8 @@ class TestSingleStageReplicaInitialization:
         runtime._init_visible_devices_baseline = "0"
 
         mocker.patch.object(runtime_mod, "inject_kv_stage_info")
-        mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value="diffusion-config")
+        od_config = SimpleNamespace(max_num_seqs=None, parallel_config=SimpleNamespace(world_size=1))
+        mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value=od_config)
         mock_register = mocker.patch(
             "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
             return_value=StageRegistrationResponse(
@@ -1174,6 +1175,7 @@ class TestSingleStageReplicaInitialization:
                 os.environ[device_env_var] = prev_device_env
 
         assert result is sentinel_client
+        assert od_config.max_num_seqs == 4
         mock_register.assert_called_once_with(
             omni_master_address="127.0.0.1",
             omni_master_port=25000,
@@ -1183,7 +1185,7 @@ class TestSingleStageReplicaInitialization:
         )
         mock_manager.assert_called_once_with(
             model="fake-model",
-            od_config="diffusion-config",
+            od_config=od_config,
             stage_init_timeout=60,
             handshake_address="tcp://127.0.0.1:26001",
             addresses=mocker.ANY,
@@ -1230,7 +1232,8 @@ class TestSingleStageReplicaInitialization:
         runtime._init_visible_devices_baseline = "0"
 
         mocker.patch.object(runtime_mod, "inject_kv_stage_info")
-        mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value="diffusion-config")
+        od_config = SimpleNamespace(max_num_seqs=None, parallel_config=SimpleNamespace(world_size=1))
+        mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value=od_config)
         mocker.patch(
             "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
             return_value=StageRegistrationResponse(

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py

@@ -37,7 +37,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.worker.request_batch import DiffusionRequestBatch
+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
 
@@ -204,6 +204,8 @@ def get_wan22_post_process_func(
         output_type: str = "np",
         sampling_params=None,
     ):
+        if sampling_params is not None and sampling_params.output_type is not None:
+            output_type = sampling_params.output_type
         if output_type == "latent":
             return video
         video_metadata = {}
@@ -233,6 +235,8 @@ def get_wan22_pre_process_func(
         prompt = request.prompt
         multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None
         raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
+        has_image = raw_image is not None and (not isinstance(raw_image, list) or bool(raw_image))
+        request.batch_compatibility_key = ("wan22_image_condition", has_image)
         if isinstance(prompt, str):
             prompt = OmniTextPrompt(prompt=prompt)
         if "additional_information" not in prompt:
@@ -288,6 +292,7 @@ class Wan22Pipeline(
     SupportsComponentDiscovery,
     WanLoraLoaderMixin,
 ):
+    supports_request_batch = True
     _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"]
     _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
     _vae_modules: ClassVar[list[str]] = ["vae"]
@@ -561,27 +566,33 @@ class Wan22Pipeline(
 
         return latents
 
-    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
+    def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]:
-        prompt: str | None = None
+        sampling_params_list = req.sampling_params_list
-        negative_prompt: str | None = None
+        common = sampling_params_list[0]
-        prompt_embeds: torch.Tensor | None = None
+        prompt_texts = [prompt if isinstance(prompt, str) else (prompt.get("prompt") or "") for prompt in req.prompts]
-        negative_prompt_embeds: torch.Tensor | None = None
+        negative_prompts = [
-        if len(req.prompts) > 1:
+            None if isinstance(prompt, str) else prompt.get("negative_prompt") for prompt in req.prompts
-            raise ValueError(
+        ]
-                """This model only supports a single prompt, not a batched request.""",
+        prompt_fields = DiffusionRequestBatch.collate_prompt_field_map(
-                """Please pass in a single prompt object or string, or a single-item list.""",
+            req.prompts,
-            )
+            {
-        if len(req.prompts) == 1:
+                "prompt_embeds": None,
-            first_prompt = req.prompts[0]
+                "negative_prompt_embeds": None,
-            prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
+            },
-            negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt")
+        )
+        prompt_embeds = prompt_fields["prompt_embeds"]
+        negative_prompt_embeds = prompt_fields["negative_prompt_embeds"]
+        prompt: list[str] | None = prompt_texts if prompt_embeds is None else None
+        negative_prompt: list[str] | None = None
+        if negative_prompt_embeds is None and any(value is not None for value in negative_prompts):
+            negative_prompt = [value or "" for value in negative_prompts]
 
-        if not prompt:
+        if prompt is not None and not all(prompt):
-            raise ValueError("Prompt is required for Wan2.2 generation.")
+            raise ValueError("Prompt is required for Wan2.2 generation when prompt_embeds are not provided.")
 
-        height = req.sampling_params.height or 480
+        height = common.height or 480
-        width = req.sampling_params.width or 832
+        width = common.width or 832
-        num_frames = req.sampling_params.num_frames or 81
+        num_frames = common.num_frames or 81
 
         # Ensure dimensions are compatible with VAE and patch size
         # For expand_timesteps mode, we need latent dims to be even (divisible by patch_size)
@@ -589,19 +600,20 @@ class Wan22Pipeline(
         mod_value = self.vae_scale_factor_spatial * patch_size[1]  # 16*2=32 for TI2V, 8*2=16 for I2V
         height = (height // mod_value) * mod_value
         width = (width // mod_value) * mod_value
-        num_steps = 40 if req.sampling_params.num_inference_steps is None else req.sampling_params.num_inference_steps
+        num_steps = 40 if common.num_inference_steps is None else common.num_inference_steps
 
-        output_type = req.sampling_params.output_type or "np"
+        output_type = common.output_type or "np"
+        num_outputs_per_prompt = common.num_outputs_per_prompt or 1
         attention_kwargs: dict | None = None
 
-        guidance_low, guidance_high = resolve_wan_guidance_scales(req.sampling_params, default_guidance_scale=4.0)
+        guidance_low, guidance_high = resolve_wan_guidance_scales(common, default_guidance_scale=4.0)
 
         # record guidance for properties
         self._guidance_scale = guidance_low
         self._guidance_scale_2 = guidance_high
 
         # Prefer engine-configured boundary_ratio, but allow per-request fallback.
-        boundary_ratio = self.boundary_ratio if self.boundary_ratio is not None else req.sampling_params.boundary_ratio
+        boundary_ratio = self.boundary_ratio if self.boundary_ratio is not None else common.boundary_ratio
 
         if boundary_ratio is None:
             boundary_ratio = 0.875
@@ -633,32 +645,49 @@ class Wan22Pipeline(
             # Fallback to text_encoder dtype if no transformer loaded
             dtype = self.text_encoder.dtype
 
-        # Seed / generator
+        generator = req.collate_request_generators(num_outputs_per_prompt, None)
-        generator = req.sampling_params.generator
+        request_latents = req.collate_request_tensors("latents", None)
-        if generator is None and req.sampling_params.seed is not None:
-            generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed)
 
         if DEBUG_PERF:
             # Sync GPU before timing to ensure accurate measurements
             current_omni_platform.synchronize()
             _t_pipeline_start = time.perf_counter()
             _t_text_enc_start = _t_pipeline_start
-        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+        do_classifier_free_guidance = guidance_low > 1.0 or guidance_high > 1.0
-            prompt=prompt,
+        if prompt_embeds is None:
-            negative_prompt=negative_prompt,
+            prompt_embeds, negative_prompt_embeds = self.encode_prompt(
-            do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0,
+                prompt=prompt,
-            num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1,
+                negative_prompt=negative_prompt,
-            max_sequence_length=req.sampling_params.max_sequence_length or 512,
+                do_classifier_free_guidance=do_classifier_free_guidance,
-            device=device,
+                num_videos_per_prompt=num_outputs_per_prompt,
-            dtype=dtype,
+                max_sequence_length=common.max_sequence_length or 512,
-        )
+                device=device,
+                dtype=dtype,
+            )
+        else:
+            prompt_embeds = prompt_embeds.to(device=device, dtype=dtype)
+            prompt_embeds = prompt_embeds.repeat_interleave(num_outputs_per_prompt, dim=0)
+            if negative_prompt_embeds is not None:
+                negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype)
+                negative_prompt_embeds = negative_prompt_embeds.repeat_interleave(num_outputs_per_prompt, dim=0)
+            elif do_classifier_free_guidance:
+                _, negative_prompt_embeds = self.encode_prompt(
+                    prompt=[""] * req.num_reqs,
+                    negative_prompt=negative_prompt,
+                    do_classifier_free_guidance=True,
+                    num_videos_per_prompt=num_outputs_per_prompt,
+                    max_sequence_length=common.max_sequence_length or 512,
+                    device=device,
+                    dtype=dtype,
+                )
 
         if DEBUG_PERF:
             current_omni_platform.synchronize()
             _t_text_enc_ms = (time.perf_counter() - _t_text_enc_start) * 1000
 
-        sample_solver = resolve_wan_sample_solver(req, default=self._sample_solver)
+        first_request = req.requests[0]
-        flow_shift = resolve_wan_flow_shift(req, self.od_config)
+        sample_solver = resolve_wan_sample_solver(first_request, default=self._sample_solver)
+        flow_shift = resolve_wan_flow_shift(first_request, self.od_config)
         if sample_solver != self._sample_solver or abs(flow_shift - self._flow_shift) > 1e-6:
             self.scheduler = build_wan_scheduler(sample_solver, flow_shift)
             self._sample_solver = sample_solver
@@ -674,37 +703,41 @@ class Wan22Pipeline(
 
         if DEBUG_PERF:
             _t_latent_prep_start = time.perf_counter()
-        multi_modal_data = req.prompts[0].get("multi_modal_data", {}) if not isinstance(req.prompts[0], str) else None
+        images: list[PIL.Image.Image | torch.Tensor | None] = []
-        raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
+        for request_prompt in req.prompts:
-        if isinstance(raw_image, list):
+            multi_modal_data = request_prompt.get("multi_modal_data", {}) if not isinstance(request_prompt, str) else {}
-            if len(raw_image) > 1:
+            raw_image = multi_modal_data.get("image")
-                logger.warning(
+            if isinstance(raw_image, list):
-                    """Received a list of image. Only a single image is supported by this model."""
+                if len(raw_image) > 1:
-                    """Taking only the first image for now."""
+                    logger.warning("Received multiple images for one Wan request; using only the first image.")
-                )
+                raw_image = raw_image[0] if raw_image else None
-            raw_image = raw_image[0]
+            if isinstance(raw_image, str):
-        if raw_image is None:
+                raw_image = PIL.Image.open(raw_image)
-            image = None
+            images.append(cast(PIL.Image.Image | torch.Tensor | None, raw_image))
-        elif isinstance(raw_image, str):
-            image = PIL.Image.open(raw_image)
-        else:
-            image = cast(PIL.Image.Image | torch.Tensor, raw_image)
 
         latent_condition = None
         first_frame_mask = None
 
-        if self.expand_timesteps and image is not None:
+        if self.expand_timesteps and any(image is not None for image in images):
+            if not all(image is not None for image in images):
+                raise ValueError("Cannot batch Wan requests with a mix of provided and missing image conditions.")
             # I2V mode: encode image and prepare condition
             from diffusers.video_processor import VideoProcessor
 
             video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
 
-            # Preprocess image
+            image_tensors = []
-            if isinstance(image, PIL.Image.Image):
+            for image in images:
-                image = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
+                assert image is not None
-                image_tensor = video_processor.preprocess(image, height=height, width=width)
+                if isinstance(image, PIL.Image.Image):
-            else:
+                    image = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
-                image_tensor = image
+                    image_tensor = video_processor.preprocess(image, height=height, width=width)
+                else:
+                    image_tensor = image.unsqueeze(0) if image.ndim == 3 else image
+                image_tensors.append(image_tensor)
+            image_tensor = DiffusionRequestBatch.collate_tensors(image_tensors, "image condition", None)
+            assert image_tensor is not None
+            image_tensor = image_tensor.repeat_interleave(num_outputs_per_prompt, dim=0)
 
             # Use out_channels for noise latents (not in_channels which includes condition)
             num_channels_latents = self.transformer_config.out_channels
@@ -720,7 +753,7 @@ class Wan22Pipeline(
                 dtype=torch.float32,
                 device=device,
                 generator=generator,
-                latents=req.sampling_params.latents,
+                latents=request_latents,
             )
 
             # Encode image condition
@@ -731,7 +764,6 @@ class Wan22Pipeline(
             image_tensor = image_tensor.unsqueeze(2)  # [B, C, 1, H, W]
             image_tensor = image_tensor.to(device=device, dtype=self.vae.dtype)
             latent_condition = retrieve_latents(self.vae.encode(image_tensor), sample_mode="argmax")
-            latent_condition = latent_condition.repeat(batch_size, 1, 1, 1, 1)
 
             # Normalize condition latents
             latents_mean = (
@@ -747,7 +779,7 @@ class Wan22Pipeline(
 
             # Create mask: 0 for first frame (condition), 1 for rest (to denoise)
             first_frame_mask = torch.ones(
-                1, 1, num_latent_frames, latent_height, latent_width, dtype=torch.float32, device=device
+                batch_size, 1, num_latent_frames, latent_height, latent_width, dtype=torch.float32, device=device
             )
             first_frame_mask[:, :, 0] = 0
         else:
@@ -762,7 +794,7 @@ class Wan22Pipeline(
                 dtype=torch.float32,
                 device=device,
                 generator=generator,
-                latents=req.sampling_params.latents,
+                latents=request_latents,
             )
         if DEBUG_PERF:
             current_omni_platform.synchronize()
@@ -839,8 +871,13 @@ class Wan22Pipeline(
                     _t_pipeline_wall_ms - _t_stages_sum,
                 )
 
-        return DiffusionOutput(
+        return split_diffusion_output_by_request(
-            output=output, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
+            DiffusionOutput(
+                output=output,
+                stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
+            ),
+            req,
+            num_outputs_per_prompt=num_outputs_per_prompt,
         )
 
     def predict_noise(

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py

@@ -45,7 +45,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.worker.request_batch import DiffusionRequestBatch
+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
 
@@ -65,6 +65,8 @@ def get_wan22_i2v_post_process_func(
         output_type: str = "np",
         sampling_params=None,
     ):
+        if sampling_params is not None and sampling_params.output_type is not None:
+            output_type = sampling_params.output_type
         if output_type == "latent":
             return video
         video_metadata = {}
@@ -84,6 +86,32 @@ def get_wan22_i2v_post_process_func(
     return post_process_func
 
 
+def _normalize_i2v_last_image(
+    value: str | PIL.Image.Image | torch.Tensor | list[object] | None,
+) -> PIL.Image.Image | torch.Tensor | None:
+    """Normalize the optional I2V last-frame condition for batching."""
+    if value is None:
+        return None
+
+    if isinstance(value, list):
+        if not value:
+            return None
+        if len(value) != 1:
+            raise ValueError("I2V accepts at most one last_image.")
+        value = value[0]
+
+    if isinstance(value, str):
+        value = PIL.Image.open(value).convert("RGB")
+
+    if not isinstance(value, (PIL.Image.Image, torch.Tensor)):
+        raise TypeError(
+            f"Unsupported last_image format {value.__class__}. "
+            "Expected a file path, PIL.Image.Image, torch.Tensor, or None."
+        )
+
+    return value
+
+
 def get_wan22_i2v_pre_process_func(
     od_config: OmniDiffusionConfig,
 ):
@@ -109,6 +137,9 @@ def get_wan22_i2v_pre_process_func(
                 """Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`""",
             )
         image = PIL.Image.open(raw_image).convert("RGB") if isinstance(raw_image, str) else raw_image
+        last_image = _normalize_i2v_last_image(multi_modal_data.get("last_image"))  # type: ignore[union-attr]
+        prompt["multi_modal_data"]["last_image"] = last_image  # type: ignore[index]
+        request.batch_compatibility_key = ("wan22_i2v_last_image", last_image is not None)
 
         # Calculate dimensions based on aspect ratio if not provided
         if request.sampling_params.height is None or request.sampling_params.width is None:
@@ -157,6 +188,7 @@ class Wan22I2VPipeline(
     Wan2.2-style I2V (with expand_timesteps for TI2V-5B).
     """
 
+    supports_request_batch = True
     _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"]
     _encoder_modules: ClassVar[list[str]] = ["text_encoder", "image_encoder"]
     _vae_modules: ClassVar[list[str]] = ["vae"]
@@ -433,60 +465,67 @@ class Wan22I2VPipeline(
         quant_config = getattr(self.od_config, "quantization_config", None)
         return create_transformer_from_config(config, quant_config=quant_config)
 
-    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
+    def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]:
-        prompt: str | None = None
+        sampling_params_list = req.sampling_params_list
-        negative_prompt: str | None = None
+        common = sampling_params_list[0]
-        prompt_embeds: torch.Tensor | None = None
+        prompt_texts = [prompt if isinstance(prompt, str) else (prompt.get("prompt") or "") for prompt in req.prompts]
-        negative_prompt_embeds: torch.Tensor | None = None
+        negative_prompts = [
+            None if isinstance(prompt, str) else prompt.get("negative_prompt") for prompt in req.prompts
+        ]
+        prompt_fields = DiffusionRequestBatch.collate_prompt_field_map(
+            req.prompts,
+            {
+                "prompt_embeds": None,
+                "negative_prompt_embeds": None,
+            },
+        )
+        prompt_embeds = prompt_fields["prompt_embeds"]
+        negative_prompt_embeds = prompt_fields["negative_prompt_embeds"]
         image_embeds: torch.Tensor | None = None
-        if len(req.prompts) > 1:
+        prompt: list[str] | None = prompt_texts if prompt_embeds is None else None
-            raise ValueError(
+        negative_prompt: list[str] | None = None
-                """This model only supports a single prompt, not a batched request.""",
+        if negative_prompt_embeds is None and any(value is not None for value in negative_prompts):
-                """Please pass in a single prompt object or string, or a single-item list.""",
+            negative_prompt = [value or "" for value in negative_prompts]
-            )
+        if prompt is not None and not all(prompt):
-        if len(req.prompts) == 1:
+            raise ValueError("Prompt is required for Wan2.2 I2V generation when prompt_embeds are not provided.")
-            first_prompt = req.prompts[0]
-            prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
-            negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt")
 
-        if not prompt:
+        images: list[PIL.Image.Image | torch.Tensor] = []
-            raise ValueError("Prompt is required for Wan2.2 generation.")
+        last_images: list[PIL.Image.Image | torch.Tensor | None] = []
+        for request_prompt in req.prompts:
+            multi_modal_data = request_prompt.get("multi_modal_data", {}) if not isinstance(request_prompt, str) else {}
+            raw_image = multi_modal_data.get("image")
+            if raw_image is None:
+                raise ValueError("Image is required for I2V generation.")
+            if isinstance(raw_image, list):
+                if len(raw_image) > 1:
+                    logger.warning("Received multiple images for one I2V request; using only the first image.")
+                raw_image = raw_image[0]
+            if isinstance(raw_image, str):
+                raw_image = PIL.Image.open(raw_image).convert("RGB")
+            images.append(cast(PIL.Image.Image | torch.Tensor, raw_image))
 
-        # Get image from request
+            last_image = multi_modal_data.get("last_image")
-        multi_modal_data = req.prompts[0].get("multi_modal_data", {}) if not isinstance(req.prompts[0], str) else None
+            if isinstance(last_image, str):
-        raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
+                last_image = PIL.Image.open(last_image).convert("RGB")
-        if raw_image is None:
+            last_images.append(cast(PIL.Image.Image | torch.Tensor | None, last_image))
-            raise ValueError("Image is required for I2V generation.")
+        if any(image is not None for image in last_images) and not all(image is not None for image in last_images):
-        if isinstance(raw_image, list):
+            raise ValueError("Cannot batch I2V requests with a mix of provided and missing last_image conditions.")
-            if len(raw_image) > 1:
-                logger.warning(
-                    """Received a list of image. Only a single image is supported by this model."""
-                    """Taking only the first image for now."""
-                )
-            raw_image = raw_image[0]
-        if isinstance(raw_image, str):
-            image = PIL.Image.open(raw_image)
-        else:
-            image = cast(PIL.Image.Image | torch.Tensor, raw_image)
 
-        last_image: PIL.Image.Image | torch.Tensor | None = None
+        height = common.height or 480
-        if multi_modal_data is not None:
+        width = common.width or 832
-            last_image = multi_modal_data.get("last_image", None)
+        num_frames = common.num_frames or 81
+        num_steps = 40 if common.num_inference_steps is None else common.num_inference_steps
 
-        height = req.sampling_params.height or 480
+        output_type = common.output_type or "np"
-        width = req.sampling_params.width or 832
+        num_outputs_per_prompt = common.num_outputs_per_prompt or 1
-        num_frames = req.sampling_params.num_frames or 81
-        num_steps = 40 if req.sampling_params.num_inference_steps is None else req.sampling_params.num_inference_steps
-
-        output_type = req.sampling_params.output_type or "np"
         attention_kwargs: dict | None = None
 
-        guidance_low, guidance_high = resolve_wan_guidance_scales(req.sampling_params, default_guidance_scale=5.0)
+        guidance_low, guidance_high = resolve_wan_guidance_scales(common, default_guidance_scale=5.0)
 
         self._guidance_scale = guidance_low
         self._guidance_scale_2 = guidance_high
 
-        boundary_ratio = self.boundary_ratio if self.boundary_ratio is not None else req.sampling_params.boundary_ratio
+        boundary_ratio = self.boundary_ratio if self.boundary_ratio is not None else common.boundary_ratio
         if boundary_ratio is None:
             boundary_ratio = 0.875
             logger.warning("boundary_ratio is required for I2V generation. using default value 0.875")
@@ -495,7 +534,7 @@ class Wan22I2VPipeline(
         self.check_inputs(
             prompt=prompt,
             negative_prompt=negative_prompt,
-            image=image,
+            image=images,
             height=height,
             width=width,
             prompt_embeds=prompt_embeds,
@@ -513,10 +552,8 @@ class Wan22I2VPipeline(
         device = self.device
         dtype = self.transformer.dtype
 
-        # Generator setup
+        generator = req.collate_request_generators(num_outputs_per_prompt, None)
-        generator = req.sampling_params.generator
+        request_latents = req.collate_request_tensors("latents", None)
-        if generator is None and req.sampling_params.seed is not None:
-            generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed)
 
         if DEBUG_PERF:
             # Sync GPU before timing to ensure accurate measurements
@@ -524,15 +561,33 @@ class Wan22I2VPipeline(
             _t_pipeline_start = time.perf_counter()
             _t_text_enc_start = _t_pipeline_start
 
-        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+        do_classifier_free_guidance = guidance_low > 1.0 or guidance_high > 1.0
-            prompt=prompt,
+        if prompt_embeds is None:
-            negative_prompt=negative_prompt,
+            prompt_embeds, negative_prompt_embeds = self.encode_prompt(
-            do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0,
+                prompt=prompt,
-            num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1,
+                negative_prompt=negative_prompt,
-            max_sequence_length=req.sampling_params.max_sequence_length or 512,
+                do_classifier_free_guidance=do_classifier_free_guidance,
-            device=device,
+                num_videos_per_prompt=num_outputs_per_prompt,
-            dtype=dtype,
+                max_sequence_length=common.max_sequence_length or 512,
-        )
+                device=device,
+                dtype=dtype,
+            )
+        else:
+            prompt_embeds = prompt_embeds.to(device=device, dtype=dtype)
+            prompt_embeds = prompt_embeds.repeat_interleave(num_outputs_per_prompt, dim=0)
+            if negative_prompt_embeds is not None:
+                negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype)
+                negative_prompt_embeds = negative_prompt_embeds.repeat_interleave(num_outputs_per_prompt, dim=0)
+            elif do_classifier_free_guidance:
+                _, negative_prompt_embeds = self.encode_prompt(
+                    prompt=[""] * req.num_reqs,
+                    negative_prompt=negative_prompt,
+                    do_classifier_free_guidance=True,
+                    num_videos_per_prompt=num_outputs_per_prompt,
+                    max_sequence_length=common.max_sequence_length or 512,
+                    device=device,
+                    dtype=dtype,
+                )
 
         if DEBUG_PERF:
             current_omni_platform.synchronize()
@@ -543,12 +598,18 @@ class Wan22I2VPipeline(
         if DEBUG_PERF:
             _t_img_enc_start = time.perf_counter()
         if self.has_image_encoder and self.transformer.config.image_dim is not None:
-            if image_embeds is None:
+            if all(last_image is None for last_image in last_images):
-                if last_image is None:
+                image_embeds = self.encode_image(images, device)
-                    image_embeds = self.encode_image(image, device)
+                image_embeds = image_embeds.repeat_interleave(num_outputs_per_prompt, dim=0)
-                else:
+            else:
-                    image_embeds = self.encode_image([image, last_image], device)
+                image_pairs = []
-            image_embeds = image_embeds.repeat(batch_size, 1, 1)
+                for first_image, last_image in zip(images, last_images):
+                    assert last_image is not None
+                    image_pairs.extend([first_image, last_image])
+                image_embeds = self.encode_image(image_pairs, device)
+                image_embeds = image_embeds.view(req.num_reqs, 2, *image_embeds.shape[1:])
+                image_embeds = image_embeds.repeat_interleave(num_outputs_per_prompt, dim=0)
+                image_embeds = image_embeds.flatten(0, 1)
             image_embeds = image_embeds.to(dtype)
         else:
             image_embeds = None
@@ -557,8 +618,9 @@ class Wan22I2VPipeline(
             current_omni_platform.synchronize()
             _t_img_enc_ms = (time.perf_counter() - _t_img_enc_start) * 1000
 
-        sample_solver = resolve_wan_sample_solver(req, default=self._sample_solver)
+        first_request = req.requests[0]
-        flow_shift = resolve_wan_flow_shift(req, self.od_config)
+        sample_solver = resolve_wan_sample_solver(first_request, default=self._sample_solver)
+        flow_shift = resolve_wan_flow_shift(first_request, self.od_config)
         if sample_solver != self._sample_solver or abs(flow_shift - self._flow_shift) > 1e-6:
             self.scheduler = build_wan_scheduler(sample_solver, flow_shift)
             self._sample_solver = sample_solver
@@ -582,23 +644,36 @@ class Wan22I2VPipeline(
 
         video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
 
-        if isinstance(image, PIL.Image.Image):
+        image_tensors = []
-            image = TF.to_tensor(image).to(device)
+        for image in images:
-            image_tensor = video_processor.preprocess(image, height=height, width=width)
+            if isinstance(image, PIL.Image.Image):
-        else:
+                image = TF.to_tensor(image).to(device)
-            image_tensor = image
+                image_tensor = video_processor.preprocess(image, height=height, width=width)
-        image_tensor = image_tensor.to(device=device, dtype=torch.float32)
-
-        # Handle last_image if provided
-        if last_image is not None:
-            if isinstance(last_image, PIL.Image.Image):
-                image = TF.to_tensor(last_image).to(device)
-                last_image_tensor = video_processor.preprocess(last_image, height=height, width=width)
             else:
-                last_image_tensor = last_image
+                image_tensor = image.unsqueeze(0) if image.ndim == 3 else image
-            last_image_tensor = last_image_tensor.to(device=device, dtype=torch.float32)
+            image_tensors.append(image_tensor.to(device=device, dtype=torch.float32))
-        else:
+        image_tensor = DiffusionRequestBatch.collate_tensors(image_tensors, "image condition", None)
-            last_image_tensor = None
+        assert image_tensor is not None
+        image_tensor = image_tensor.repeat_interleave(num_outputs_per_prompt, dim=0)
+
+        last_image_tensor = None
+        if all(last_image is not None for last_image in last_images):
+            last_image_tensors = []
+            for last_image in last_images:
+                assert last_image is not None
+                if isinstance(last_image, PIL.Image.Image):
+                    last_image = TF.to_tensor(last_image).to(device)
+                    current_last_image = video_processor.preprocess(last_image, height=height, width=width)
+                else:
+                    current_last_image = last_image.unsqueeze(0) if last_image.ndim == 3 else last_image
+                last_image_tensors.append(current_last_image.to(device=device, dtype=torch.float32))
+            last_image_tensor = DiffusionRequestBatch.collate_tensors(
+                last_image_tensors,
+                "last image condition",
+                None,
+            )
+            assert last_image_tensor is not None
+            last_image_tensor = last_image_tensor.repeat_interleave(num_outputs_per_prompt, dim=0)
 
         latents, condition, first_frame_mask = self.prepare_latents(
             image=image_tensor,
@@ -610,7 +685,7 @@ class Wan22I2VPipeline(
             dtype=torch.float32,
             device=device,
             generator=generator,
-            latents=req.sampling_params.latents,
+            latents=request_latents,
             last_image=last_image_tensor,
         )
 
@@ -694,8 +769,13 @@ class Wan22I2VPipeline(
                     _t_pipeline_wall_ms - _t_stages_sum,
                 )
 
-        return DiffusionOutput(
+        return split_diffusion_output_by_request(
-            output=output, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
+            DiffusionOutput(
+                output=output,
+                stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
+            ),
+            req,
+            num_outputs_per_prompt=num_outputs_per_prompt,
         )
 
     def predict_noise(
@@ -849,7 +929,6 @@ class Wan22I2VPipeline(
 
         # Encode through VAE
         latent_condition = retrieve_latents(self.vae.encode(video_condition), sample_mode="argmax")
-        latent_condition = latent_condition.repeat(batch_size, 1, 1, 1, 1)
 
         # Normalize latents
         latents_mean = (
@@ -866,7 +945,7 @@ class Wan22I2VPipeline(
         if self.expand_timesteps:
             # TI2V-5B style: create mask where first frame is 0 (condition), rest is 1 (to denoise)
             first_frame_mask = torch.ones(
-                1, 1, num_latent_frames, latent_height, latent_width, dtype=dtype, device=device
+                batch_size, 1, num_latent_frames, latent_height, latent_width, dtype=dtype, device=device
             )
             first_frame_mask[:, :, 0] = 0
             return latents, latent_condition, first_frame_mask
@@ -892,7 +971,9 @@ class Wan22I2VPipeline(
         condition = torch.concat([mask_lat_size, latent_condition], dim=1)
 
         # For non-expand mode, first_frame_mask is not used in the same way
-        first_frame_mask = torch.ones(1, 1, num_latent_frames, latent_height, latent_width, dtype=dtype, device=device)
+        first_frame_mask = torch.ones(
+            batch_size, 1, num_latent_frames, latent_height, latent_width, dtype=dtype, device=device
+        )
 
         return latents, condition, first_frame_mask
 

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py

@@ -43,7 +43,7 @@ from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import (
 )
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
-from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+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 +63,30 @@ _S2V_DEFAULT_NEG_PROMPT = (
 _AUDIO_VIDEO_RATE = 30
 
 
+def _s2v_audio_condition_key(audio: str | np.ndarray | None) -> tuple[Any, ...]:
+    if isinstance(audio, str):
+        # File paths are deliberately conservative: two different files may
+        # have different durations even when their request metadata matches.
+        return ("path", audio)
+    if isinstance(audio, np.ndarray):
+        return ("array", tuple(audio.shape), str(audio.dtype))
+    return (type(audio).__name__,)
+
+
+def _make_clip_generators(
+    seeds: list[int | None],
+    request_generators: list[torch.Generator | None],
+    clip_index: int,
+    device: torch.device,
+) -> list[torch.Generator | None]:
+    return [
+        torch.Generator(device=device).manual_seed((seed if seed is not None else 0) + clip_index)
+        if seed is not None or request_generator is None
+        else request_generator
+        for seed, request_generator in zip(seeds, request_generators)
+    ]
+
+
 def _linear_interpolation(
     features: torch.Tensor,
     input_fps: float,
@@ -337,6 +361,12 @@ def get_wan22_s2v_pre_process_func(
         prompt["additional_information"]["init_first_frame"] = (
             multi_modal_data.get("init_first_frame", False) if multi_modal_data is not None else False
         )
+        request.batch_compatibility_key = (
+            "wan22_s2v_condition",
+            bool(prompt["additional_information"]["init_first_frame"]),
+            prompt["additional_information"].get("num_repeat"),
+            _s2v_audio_condition_key(raw_audio),
+        )
         request.prompt = prompt
 
         return request
@@ -486,6 +516,7 @@ class Wan22S2VPipeline(
     """
 
     # Default config values from Wan2.2/wan/configs/wan_s2v_14B.py
+    supports_request_batch = True
     _DEFAULT_MOTION_FRAMES = 73
     _DEFAULT_INFER_FRAMES = 80
     _DEFAULT_FPS = 16
@@ -942,7 +973,6 @@ class Wan22S2VPipeline(
         """Forward pass through the S2V transformer to predict noise.
 
         With return_dict=False, the model returns (output,) where output is [B, C, T, H, W].
-        We squeeze batch dim since S2V processes one sample at a time.
         """
         if current_model is None:
             current_model = self.transformer
@@ -952,8 +982,8 @@ class Wan22S2VPipeline(
         if isinstance(result, list):
             return result[0]
         if isinstance(result, tuple):
-            return result[0].squeeze(0)
+            return result[0]
-        return result.sample.squeeze(0)
+        return result.sample
 
     def diffuse(
         self,
@@ -962,7 +992,7 @@ class Wan22S2VPipeline(
         prompt_embeds: torch.Tensor,
         negative_prompt_embeds: torch.Tensor | None,
         guidance_scale: float,
-        clip_generator: torch.Generator,
+        clip_generator: torch.Generator | list[torch.Generator] | None,
         dtype: torch.dtype,
         device: torch.device,
         max_seq_len: int,
@@ -971,13 +1001,13 @@ class Wan22S2VPipeline(
         ref_latents: torch.Tensor,
         motion_frames: list[int],
         drop_first_motion: bool,
-        positive_audio_emb: torch.Tensor,
+        positive_audio_emb: dict[str, torch.Tensor],
-        negative_audio_emb: torch.Tensor | None,
+        negative_audio_emb: dict[str, torch.Tensor] | None,
     ) -> torch.Tensor:
-        """Denoising diffusion loop for one S2V clip.
+        """Denoising diffusion loop for a batch of compatible S2V clips.
 
         Args:
-            latents: Initial noise tensor [C, T, H, W]
+            latents: Initial noise tensor [B, C, T, H, W]
             timesteps: Denoising timesteps
             prompt_embeds: Text embeddings [1, seq_len, dim]
             negative_prompt_embeds: Negative text embeddings (optional)
@@ -995,7 +1025,7 @@ class Wan22S2VPipeline(
             negative_audio_emb: Precomputed audio embeddings for negative prompt (optional)
 
         Returns:
-            Denoised latents [C, T, H, W]
+            Denoised latents [B, C, T, H, W]
         """
         do_true_cfg = self.do_classifier_free_guidance and negative_prompt_embeds is not None
 
@@ -1004,13 +1034,13 @@ class Wan22S2VPipeline(
                 self._current_timestep = t
                 self.record_denoise_step(step_idx, t)
 
-                latent_model_input = [latents.to(device)]
+                latent_model_input = latents.to(device)
-                timestep = t.unsqueeze(0).to(device) if t.dim() == 0 else t.to(device)
+                timestep = t.expand(latents.shape[0]).to(device)
 
                 positive_kwargs = {
                     "hidden_states": latent_model_input,
                     "timestep": timestep,
-                    "encoder_hidden_states": prompt_embeds[0:1],
+                    "encoder_hidden_states": prompt_embeds,
                     "cond_states": cond_latents,
                     "motion_latents": input_motion_latents,
                     "ref_latents": ref_latents,
@@ -1023,7 +1053,7 @@ class Wan22S2VPipeline(
                     negative_kwargs = {
                         "hidden_states": latent_model_input,
                         "timestep": timestep,
-                        "encoder_hidden_states": negative_prompt_embeds[0:1],
+                        "encoder_hidden_states": negative_prompt_embeds,
                         "cond_states": cond_latents,
                         "motion_latents": input_motion_latents,
                         "ref_latents": ref_latents,
@@ -1044,12 +1074,12 @@ class Wan22S2VPipeline(
 
                 # Scheduler step
                 latents = self.scheduler.step(
-                    noise_pred.unsqueeze(0),
+                    noise_pred,
                     t,
-                    latents.unsqueeze(0),
+                    latents,
                     return_dict=False,
-                    generator=clip_generator,
+                    generator=clip_generator if isinstance(clip_generator, torch.Generator) else None,
-                )[0].squeeze(0)
+                )[0]
 
                 pbar.update()
 
@@ -1080,111 +1110,178 @@ class Wan22S2VPipeline(
         prompt_embeds: torch.Tensor | None = None,
         negative_prompt_embeds: torch.Tensor | None = None,
         **kwargs,
-    ) -> DiffusionOutput:
+    ) -> list[DiffusionOutput]:
         """Run S2V generation — may produce multiple autoregressive clips.
 
         This method mirrors ``WanS2V.generate()``, reorganized into the
         vLLM-Omni pipeline structure.
         """
-        # ---- Extract params from request ----
+        sampling_params_list = req.sampling_params_list
-        if len(req.prompts) > 1:
+        common = sampling_params_list[0]
-            raise ValueError("S2V only supports a single prompt per request.")
+        prompt_fields = DiffusionRequestBatch.collate_prompt_field_map(
-
+            req.prompts,
-        if len(req.prompts) == 1:
+            {
-            prompt_data = req.prompts[0]
+                "prompt_embeds": prompt_embeds,
-            prompt = prompt_data if isinstance(prompt_data, str) else prompt_data.get("prompt")
+                "negative_prompt_embeds": negative_prompt_embeds,
-            negative_prompt = None if isinstance(prompt_data, str) else prompt_data.get("negative_prompt")
+            },
+        )
+        prompt_embeds = prompt_fields["prompt_embeds"]
+        negative_prompt_embeds = prompt_fields["negative_prompt_embeds"]
 
+        prompt_texts: list[str] = []
+        negative_prompts: list[str] = []
+        images: list[PIL.Image.Image | None] = []
+        audio_paths: list[str | np.ndarray | None] = []
+        requested_repeats: list[int | None] = []
+        init_first_frames: list[bool] = []
+        for prompt_data in req.prompts:
             multi_modal_data = prompt_data.get("multi_modal_data", {}) if not isinstance(prompt_data, str) else {}
             additional_info = prompt_data.get("additional_information", {}) if not isinstance(prompt_data, str) else {}
 
-            if image is None:
+            prompt_texts.append(
-                raw_image = multi_modal_data.get("image", None)
+                prompt_data if isinstance(prompt_data, str) else (prompt_data.get("prompt") or prompt or "")
-                if isinstance(raw_image, str):
+            )
-                    image = PIL.Image.open(raw_image).convert("RGB")
+            request_negative_prompt = (
-                elif isinstance(raw_image, PIL.Image.Image):
+                negative_prompt if isinstance(prompt_data, str) else prompt_data.get("negative_prompt", negative_prompt)
-                    image = raw_image
+            )
+            negative_prompts.append(request_negative_prompt or _S2V_DEFAULT_NEG_PROMPT)
 
-            if audio_path is None:
+            raw_image = multi_modal_data.get("image", image)
-                audio_path = additional_info.get("audio_path")
+            if isinstance(raw_image, list):
-                if audio_path is None:
+                raw_image = raw_image[0] if raw_image else None
-                    audio_path = multi_modal_data.get("audio")
+            if isinstance(raw_image, str):
+                raw_image = PIL.Image.open(raw_image).convert("RGB")
+            images.append(raw_image if isinstance(raw_image, PIL.Image.Image) else None)
 
-            if pose_video is None:
+            request_audio_path = additional_info.get("audio_path", audio_path)
-                pose_video = additional_info.get("pose_video")
+            if request_audio_path is None:
+                request_audio_path = multi_modal_data.get("audio")
+            audio_paths.append(request_audio_path)
+            requested_repeats.append(additional_info.get("num_repeat", num_repeat))
+            init_first_frames.append(additional_info.get("init_first_frame", init_first_frame))
 
-            init_first_frame = additional_info.get("init_first_frame", init_first_frame)
+        height = common.height or height
-
+        width = common.width or width
-        if negative_prompt is None:
-            negative_prompt = _S2V_DEFAULT_NEG_PROMPT
-
-        height = req.sampling_params.height or height
-        width = req.sampling_params.width or width
         # num_frames defaults to 1 (image-model sentinel); S2V needs many frames.
-        req_num_frames = req.sampling_params.num_frames
+        req_num_frames = common.num_frames
         infer_frames = (
             infer_frames
             or (req_num_frames if req_num_frames and req_num_frames > 1 else None)
             or self._DEFAULT_INFER_FRAMES
         )
-        num_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        num_steps = common.num_inference_steps or num_inference_steps
+        num_outputs_per_prompt = common.num_outputs_per_prompt or 1
 
-        if req.sampling_params.guidance_scale_provided:
+        if common.guidance_scale_provided:
-            guidance_scale = req.sampling_params.guidance_scale
+            guidance_scale = common.guidance_scale
 
         self._guidance_scale = guidance_scale
 
         # ---- Validate ----
-        self.check_inputs(prompt, image, audio_path, height, width, prompt_embeds)
+        if len(set(init_first_frames)) != 1:
+            raise ValueError("Batched S2V requests must have matching init_first_frame values.")
+        init_first_frame = init_first_frames[0]
+        for request_prompt, request_image, request_audio_path in zip(prompt_texts, images, audio_paths):
+            self.check_inputs(
+                request_prompt if prompt_embeds is None else None,
+                request_image,
+                request_audio_path,
+                height,
+                width,
+                prompt_embeds,
+            )
 
         device = self.device
         dtype = self.transformer.dtype if hasattr(self.transformer, "dtype") else torch.bfloat16
-
+        request_generators = req.collate_request_generators(num_outputs_per_prompt, None)
-        if generator is None:
+        batch_size = req.num_reqs * num_outputs_per_prompt
-            generator = req.sampling_params.generator
+        if isinstance(request_generators, list):
-        seed = req.sampling_params.seed
+            generators = request_generators
-        if generator is None and seed is not None:
+        else:
-            generator = torch.Generator(device=device).manual_seed(seed)
+            generators = [request_generators] * batch_size
-        if seed is None:
+        seeds = [sampling.seed for sampling in sampling_params_list for _ in range(num_outputs_per_prompt)]
-            seed = 0
+        request_latents = req.collate_request_tensors("latents", None)
 
         # ---- 1. Text encoding ----
         if prompt_embeds is None:
             prompt_embeds, negative_prompt_embeds = self.encode_prompt(
-                prompt=prompt,
+                prompt=prompt_texts,
-                negative_prompt=negative_prompt,
+                negative_prompt=negative_prompts,
                 do_classifier_free_guidance=self.do_classifier_free_guidance,
-                num_videos_per_prompt=1,
+                num_videos_per_prompt=num_outputs_per_prompt,
-                max_sequence_length=512,
+                max_sequence_length=common.max_sequence_length or 512,
                 device=device,
                 dtype=dtype,
             )
         else:
             prompt_embeds = prompt_embeds.to(device=device, dtype=dtype)
+            prompt_embeds = prompt_embeds.repeat_interleave(num_outputs_per_prompt, dim=0)
             if negative_prompt_embeds is not None:
                 negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype)
+                negative_prompt_embeds = negative_prompt_embeds.repeat_interleave(num_outputs_per_prompt, dim=0)
 
         # ---- 2. Audio encoding ----
-        audio_emb, audio_num_repeat, target_video_frames = self.encode_audio(
+        audio_embeddings: list[torch.Tensor] = []
-            audio_path, infer_frames=infer_frames, device=device, dtype=dtype
+        effective_repeats: list[int] = []
-        )
+        target_video_frames: list[int] = []
-        if num_repeat is None or num_repeat > audio_num_repeat:
+        for request_audio_path, requested_repeat in zip(audio_paths, requested_repeats):
-            num_repeat = audio_num_repeat
+            assert request_audio_path is not None
+            audio_emb, audio_num_repeat, request_target_video_frames = self.encode_audio(
+                request_audio_path, infer_frames=infer_frames, device=device, dtype=dtype
+            )
+            audio_embeddings.append(audio_emb.repeat_interleave(num_outputs_per_prompt, dim=0))
+            effective_repeats.append(
+                audio_num_repeat
+                if requested_repeat is None or requested_repeat > audio_num_repeat
+                else requested_repeat
+            )
+            target_video_frames.append(request_target_video_frames)
+
+        if len(set(effective_repeats)) != 1:
+            raise ValueError(f"Batched S2V requests must have matching num_repeat values, got {effective_repeats}.")
+        if len(set(target_video_frames)) != 1:
+            raise ValueError(
+                "Batched S2V requests must have matching audio embedding time lengths; "
+                f"got target video frames {target_video_frames}."
+            )
+        num_repeat = effective_repeats[0]
+        audio_emb = DiffusionRequestBatch.collate_tensors(audio_embeddings, "S2V audio embeddings", None)
+        assert audio_emb is not None
 
         # Load raw audio waveform for muxing into output video
-        if isinstance(audio_path, np.ndarray):
+        raw_audio_waveforms: list[np.ndarray] = []
-            raw_audio_waveform = audio_path.astype(np.float32)
+        raw_audio_sample_rates: list[int] = []
-            raw_audio_sr = 16000
+        for request_audio_path in audio_paths:
-        else:
+            assert request_audio_path is not None
-            raw_audio_waveform, raw_audio_sr = load_audio(audio_path, sr=None, mono=True)
+            if isinstance(request_audio_path, np.ndarray):
+                raw_audio_waveforms.append(request_audio_path.astype(np.float32, copy=False))
+                raw_audio_sample_rates.append(16000)
+            else:
+                raw_audio_waveform, raw_audio_sr = load_audio(request_audio_path, sr=None, mono=True)
+                raw_audio_waveforms.append(raw_audio_waveform)
+                raw_audio_sample_rates.append(raw_audio_sr)
+        if len(set(raw_audio_sample_rates)) != 1 or len({waveform.shape for waveform in raw_audio_waveforms}) != 1:
+            raise ValueError("Batched S2V requests must have matching raw audio shapes and sample rates.")
+        raw_audio_sr = raw_audio_sample_rates[0]
 
         # ---- 3. Reference image encoding ----
-        ref_latents = self.encode_ref_image(image, height, width, device=device).to(dtype=dtype)
+        ref_latents = DiffusionRequestBatch.collate_tensors(
+            [
+                self.encode_ref_image(request_image, height, width, device=device)
+                .to(dtype=dtype)
+                .repeat_interleave(num_outputs_per_prompt, dim=0)
+                for request_image in images
+                if request_image is not None
+            ],
+            "S2V reference image latents",
+            None,
+        )
+        assert ref_latents is not None
 
         # ---- 4. Initial motion latents (zeros) ----
         motion_frames = min(self.motion_frames, infer_frames - 1)
         motion_pixels = torch.zeros(
-            [1, 3, motion_frames, height, width],
+            [batch_size, 3, motion_frames, height, width],
             dtype=dtype,
             device=device,
         )
@@ -1193,8 +1290,11 @@ class Wan22S2VPipeline(
             drop_first_motion = False
             # Place reference image in the last 6 frames of motion
             tensor_trans = transforms.ToTensor()
-            ref_tensor = tensor_trans(image).to(device=device, dtype=dtype) * 2 - 1.0  # [C, H, W]
+            ref_tensors = torch.stack(
-            motion_pixels[:, :, -6:] = ref_tensor.unsqueeze(0).unsqueeze(2).expand(-1, -1, 6, -1, -1)
+                [tensor_trans(request_image) for request_image in images if request_image is not None]
+            ).to(device=device, dtype=dtype)
+            ref_tensors = ref_tensors.repeat_interleave(num_outputs_per_prompt, dim=0) * 2 - 1.0
+            motion_pixels[:, :, -6:] = ref_tensors.unsqueeze(2).expand(-1, -1, 6, -1, -1)
 
         motion_latents = self.prepare_motion_latents(motion_pixels, device=device).to(dtype=dtype)
 
@@ -1211,24 +1311,29 @@ class Wan22S2VPipeline(
         clips: list[torch.Tensor] = []
         # Keep a pixel-space buffer of the trailing motion_frames for the
         # autoregressive connection between clips.
-        videos_last_frames = torch.zeros([1, 3, motion_frames, height, width], dtype=dtype, device=device)
+        videos_last_frames = torch.zeros([batch_size, 3, motion_frames, height, width], dtype=dtype, device=device)
 
         for r in range(num_repeat):
-            # Per-clip seed
+            clip_generators = _make_clip_generators(seeds, generators, r, device)
-            clip_seed = seed + r
-            clip_generator = torch.Generator(device=device).manual_seed(clip_seed)
-
             # -- Noise --
-            latents = self.prepare_latents(
+            if r == 0 and request_latents is not None:
-                infer_frames=infer_frames,
+                latents = request_latents.to(device=device, dtype=dtype)
-                height=height,
+            else:
-                width=width,
+                latents = torch.stack(
-                motion_frames=motion_frames,
+                    [
-                lat_motion_frames=lat_motion_frames,
+                        self.prepare_latents(
-                dtype=dtype,
+                            infer_frames=infer_frames,
-                device=device,
+                            height=height,
-                generator=clip_generator,
+                            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)
@@ -1242,11 +1347,11 @@ class Wan22S2VPipeline(
 
             # -- Pose condition for this clip --
             # Default: zero condition (no pose driving)
-            # Shape: [B=1, C=16, T, H, W] — passed to transformer which iterates over batch dim
+            # 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([1, 16, lat_target_frames, lat_h, lat_w], dtype=dtype, device=device)
+            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()
@@ -1282,7 +1387,7 @@ class Wan22S2VPipeline(
                 prompt_embeds=prompt_embeds,
                 negative_prompt_embeds=negative_prompt_embeds,
                 guidance_scale=guidance_scale,
-                clip_generator=clip_generator,
+                clip_generator=clip_generators,
                 dtype=dtype,
                 device=device,
                 max_seq_len=max_seq_len,
@@ -1300,11 +1405,10 @@ class Wan22S2VPipeline(
                 self.transformer.to("cpu")
                 current_omni_platform.empty_cache()
 
-            latents_for_decode = latents.unsqueeze(0)  # [1, C, T, H, W]
             if not (drop_first_motion and r == 0):
-                decode_latents = torch.cat([motion_latents, latents_for_decode], dim=2)
+                decode_latents = torch.cat([motion_latents, latents], dim=2)
             else:
-                decode_latents = torch.cat([ref_latents, latents_for_decode], dim=2)
+                decode_latents = torch.cat([ref_latents, latents], dim=2)
 
             decode_latents = self._denormalize_latents(decode_latents)
             decode_latents = decode_latents.to(self.vae.dtype)
@@ -1321,7 +1425,7 @@ class Wan22S2VPipeline(
 
                 # Create buffer for broadcast
                 clip_video = torch.empty(
-                    (1, 3, total_frames, height, width),
+                    (batch_size, 3, total_frames, height, width),
                     device=decode_latents.device,
                     dtype=decode_latents.dtype,
                 )
@@ -1356,12 +1460,19 @@ class Wan22S2VPipeline(
                 current_omni_platform.empty_cache()
 
         # ---- Concatenate all clips ----
-        output = torch.cat(clips, dim=2)  # [1, C, T_total, H, W]
+        output = torch.cat(clips, dim=2)  # [B, C, T_total, H, W]
 
-        return DiffusionOutput(
+        outputs = split_diffusion_output_by_request(
-            output=(output, raw_audio_waveform, raw_audio_sr),
+            DiffusionOutput(
-            stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
+                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
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         """Load weights using AutoWeightsLoader for vLLM integration."""

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py

@@ -35,13 +35,27 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
 )
 from vllm_omni.diffusion.models.wan2_2.wan2_2_vace_transformer import WanVACETransformer3DModel
 from vllm_omni.diffusion.request import OmniDiffusionRequest
-from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+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
 
 logger = init_logger(__name__)
 
 
+def _vace_condition_structure(value: object) -> tuple:
+    if value is None:
+        return ("none",)
+    if isinstance(value, str):
+        return ("path", value)
+    if isinstance(value, PIL.Image.Image):
+        return ("image", value.size)
+    if isinstance(value, torch.Tensor):
+        return ("tensor", tuple(value.shape), str(value.dtype))
+    if isinstance(value, list):
+        return ("list", tuple(_vace_condition_structure(item) for item in value))
+    return (type(value).__name__,)
+
+
 def create_vace_transformer_from_config(
     config: dict,
     quant_config: QuantizationConfig | None = None,
@@ -104,6 +118,7 @@ def get_wan22_vace_pre_process_func(od_config: OmniDiffusionConfig):
             prompt["additional_information"] = {}
 
         if not multi_modal_data:
+            request.batch_compatibility_key = ("wan22_vace_condition", ("none",), ("none",), ("none",))
             request.prompt = prompt
             return request
 
@@ -155,6 +170,13 @@ def get_wan22_vace_pre_process_func(od_config: OmniDiffusionConfig):
                 mask = [mask]
             prompt["additional_information"]["mask"] = mask
 
+        request.batch_compatibility_key = (
+            "wan22_vace_condition",
+            _vace_condition_structure(ref_images),
+            _vace_condition_structure(source_video),
+            _vace_condition_structure(mask),
+        )
+
         request.prompt = prompt
         return request
 
@@ -482,7 +504,7 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
         attention_kwargs: dict | None = None,
         vace_context_scale: float | list[float] = 1.0,
         **kwargs,
-    ) -> DiffusionOutput:
+    ) -> list[DiffusionOutput]:
         """Generate or edit video using VACE.
 
         The mode is determined by which inputs are provided in the request:
@@ -507,47 +529,48 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
             attention_kwargs: Additional kwargs for attention layers.
             vace_context_scale: VACE conditioning strength.
         """
-        # Get parameters from request or arguments
+        sampling_params_list = req.sampling_params_list
-        if len(req.prompts) > 1:
+        common = sampling_params_list[0]
-            raise ValueError(
+        prompt_texts = [item if isinstance(item, str) else (item.get("prompt") or "") for item in req.prompts]
-                "This model only supports a single prompt, not a batched request. "
+        negative_prompts = [None if isinstance(item, str) else item.get("negative_prompt") for item in req.prompts]
-                "Please pass in a single prompt object or string, or a single-item list."
+        prompt_fields = DiffusionRequestBatch.collate_prompt_field_map(
-            )
+            req.prompts,
-
+            {
-        reference_images = None
+                "prompt_embeds": prompt_embeds,
-        source_video = None
+                "negative_prompt_embeds": negative_prompt_embeds,
-        source_mask = None
+            },
-
+        )
-        if len(req.prompts) == 1:
+        prompt_embeds = prompt_fields["prompt_embeds"]
-            first_prompt = req.prompts[0]
+        negative_prompt_embeds = prompt_fields["negative_prompt_embeds"]
-            if isinstance(first_prompt, str):
+        prompt = prompt_texts if prompt_embeds is None else None
-                prompt = first_prompt
+        negative_prompt = None
-            else:
+        if negative_prompt_embeds is None and any(value is not None for value in negative_prompts):
-                prompt = first_prompt.get("prompt")
+            negative_prompt = [value or "" for value in negative_prompts]
-                negative_prompt = negative_prompt or first_prompt.get("negative_prompt")
+        if prompt is not None and not all(prompt):
-                prompt_embeds = prompt_embeds if prompt_embeds is not None else first_prompt.get("prompt_embeds")
-                negative_prompt_embeds = (
-                    negative_prompt_embeds
-                    if negative_prompt_embeds is not None
-                    else first_prompt.get("negative_prompt_embeds")
-                )
-
-                additional_info = first_prompt.get("additional_information", {})
-                reference_images = additional_info.get("reference_images")
-                source_video = additional_info.get("source_video")
-                source_mask = additional_info.get("mask")
-
-        if prompt is None and prompt_embeds is None:
             raise ValueError("Prompt or prompt_embeds is required for VACE generation.")
 
-        height = req.sampling_params.height or height
+        reference_images_list = []
-        width = req.sampling_params.width or width
+        source_videos = []
-        num_frames = req.sampling_params.num_frames or frame_num
+        source_masks = []
-        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        for request_prompt in req.prompts:
-        generator = req.sampling_params.generator or generator
+            additional_info = (
+                request_prompt.get("additional_information", {}) if not isinstance(request_prompt, str) else {}
+            )
+            reference_images_list.append(additional_info.get("reference_images"))
+            source_videos.append(additional_info.get("source_video"))
+            source_masks.append(additional_info.get("mask"))
 
-        if req.sampling_params.guidance_scale_provided:
+        height = common.height or height
-            guidance_scale = req.sampling_params.guidance_scale
+        width = common.width or width
+        num_frames = common.num_frames or frame_num
+        num_inference_steps = common.num_inference_steps or num_inference_steps
+        num_outputs_per_prompt = common.num_outputs_per_prompt or 1
+        generator = req.collate_request_generators(num_outputs_per_prompt, None)
+        request_latents = req.collate_request_tensors("latents", None)
+        output_type = common.output_type or output_type
+
+        if common.guidance_scale_provided:
+            guidance_scale = common.guidance_scale
 
         # Ensure dimensions are compatible with VAE and patch size
         mod_value = self.vae_scale_factor_spatial * 2  # 8 * 2 = 16
@@ -558,47 +581,47 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
             num_frames = num_frames // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1
         num_frames = max(num_frames, 1)
 
-        self.check_inputs(
+        for index in range(req.num_reqs):
-            prompt=prompt,
+            self.check_inputs(
-            negative_prompt=negative_prompt,
+                prompt=None if prompt is None else prompt[index],
-            height=height,
+                negative_prompt=None if negative_prompt is None else negative_prompt[index],
-            width=width,
+                height=height,
-            prompt_embeds=prompt_embeds,
+                width=width,
-            negative_prompt_embeds=negative_prompt_embeds,
+                prompt_embeds=None if prompt_embeds is None else prompt_embeds[index],
-            video=source_video,
+                negative_prompt_embeds=None if negative_prompt_embeds is None else negative_prompt_embeds[index],
-            mask=source_mask,
+                video=source_videos[index],
-            reference_images=reference_images,
+                mask=source_masks[index],
-        )
+                reference_images=reference_images_list[index],
+            )
 
         device = self.device
         self._guidance_scale = guidance_scale
         dtype = self.transformer.dtype if self.transformer is not None else torch.bfloat16
 
-        if generator is None and req.sampling_params.seed is not None:
-            generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed)
-
         # Encode prompts
         if prompt_embeds is None:
-            if prompt is None:
-                raise ValueError("Either prompt or prompt_embeds must be provided.")
             prompt_embeds, negative_prompt_embeds = self.encode_prompt(
                 prompt=prompt,
                 negative_prompt=negative_prompt,
                 do_classifier_free_guidance=guidance_scale > 1.0,
-                num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1,
+                num_videos_per_prompt=num_outputs_per_prompt,
-                max_sequence_length=req.sampling_params.max_sequence_length or 512,
+                max_sequence_length=common.max_sequence_length or 512,
                 device=device,
                 dtype=dtype,
             )
         else:
             prompt_embeds = prompt_embeds.to(device=device, dtype=dtype)
+            prompt_embeds = prompt_embeds.repeat_interleave(num_outputs_per_prompt, dim=0)
             if negative_prompt_embeds is not None:
                 negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype)
+                negative_prompt_embeds = negative_prompt_embeds.repeat_interleave(num_outputs_per_prompt, dim=0)
             elif guidance_scale > 1.0:
                 _, negative_prompt_embeds = self.encode_prompt(
-                    prompt="",
+                    prompt=[""] * req.num_reqs,
-                    negative_prompt=None,
+                    negative_prompt=negative_prompt,
                     do_classifier_free_guidance=True,
+                    num_videos_per_prompt=num_outputs_per_prompt,
+                    max_sequence_length=common.max_sequence_length or 512,
                     device=device,
                     dtype=dtype,
                 )
@@ -608,24 +631,43 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
         # so read it from whichever transformer exists.
         active_transformer = self.transformer if self.transformer is not None else self.transformer_2
         if active_transformer.vace_patch_embedding is not None:
-            video, mask, ref_images_processed = self.preprocess_conditions(
+            contexts = []
-                video=source_video,
+            reference_counts = []
-                mask=source_mask,
+            for index, (source_video, source_mask, reference_images) in enumerate(
-                reference_images=reference_images,
+                zip(source_videos, source_masks, reference_images_list)
-                height=height,
+            ):
-                width=width,
+                video, mask, ref_images_processed = self.preprocess_conditions(
-                num_frames=num_frames,
+                    video=source_video,
-                dtype=dtype,
+                    mask=source_mask,
-                device=device,
+                    reference_images=reference_images,
-            )
+                    height=height,
+                    width=width,
+                    num_frames=num_frames,
+                    dtype=dtype,
+                    device=device,
+                )
+                request_generator = sampling_params_list[index].generator
+                if isinstance(request_generator, list):
+                    request_generator = request_generator[0]
+                conditioning_latents = self.prepare_video_latents(
+                    video,
+                    mask,
+                    ref_images_processed,
+                    request_generator,
+                    device,
+                )
+                mask_encoded = self.prepare_masks(mask, ref_images_processed)
+                context = torch.cat([conditioning_latents, mask_encoded], dim=1)
+                contexts.append(context.repeat_interleave(num_outputs_per_prompt, dim=0))
+                reference_counts.append(len(ref_images_processed[0]) if ref_images_processed else 0)
 
-            conditioning_latents = self.prepare_video_latents(video, mask, ref_images_processed, generator, device)
+            if len(set(reference_counts)) != 1:
-            mask_encoded = self.prepare_masks(mask, ref_images_processed)
+                raise ValueError(
-
+                    f"Batched VACE conditions must have the same number of reference images; got {reference_counts}."
-            # Unified VACE context: [video_latents, mask] along channels -> [B, C, T, H, W]
+                )
-            vace_context = torch.cat([conditioning_latents, mask_encoded], dim=1)
+            vace_context = DiffusionRequestBatch.collate_tensors(contexts, "VACE condition", None)
-
+            assert vace_context is not None
-            num_reference_images = len(ref_images_processed[0]) if ref_images_processed else 0
+            num_reference_images = reference_counts[0]
         else:
             vace_context = None
 
@@ -641,7 +683,7 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
             dtype=torch.float32,
             device=device,
             generator=generator,
-            latents=req.sampling_params.latents,
+            latents=request_latents,
         )
 
         # Set up scheduler
@@ -654,9 +696,7 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
         # stays None and diffuse() uses the same transformer for every step.
         boundary_timestep = None
         if self.transformer_2 is not None:
-            boundary_ratio = (
+            boundary_ratio = self.boundary_ratio if self.boundary_ratio is not None else common.boundary_ratio
-                self.boundary_ratio if self.boundary_ratio is not None else req.sampling_params.boundary_ratio
-            )
             if boundary_ratio is None:
                 boundary_ratio = 0.875
                 logger.warning_once("boundary_ratio is required for Wan2.2 VACE generation. using default value 0.875")
@@ -700,7 +740,11 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
             latents = latents / latents_std + latents_mean
             output = self.vae.decode(latents, return_dict=False)[0]
 
-        return DiffusionOutput(output=output)
+        return split_diffusion_output_by_request(
+            DiffusionOutput(output=output),
+            req,
+            num_outputs_per_prompt=num_outputs_per_prompt,
+        )
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         """Load weights using AutoWeightsLoader for vLLM integration."""

vllm_omni/diffusion/request.py

@@ -43,6 +43,10 @@ class OmniDiffusionRequest:
     # BaseScheduler takes ownership and clears this field before the request is
     # sent to a Worker.
     diffusion_kv_requests: tuple[DiffusionKVRequest, ...] | None = None
+    # Optional model-specific condition structure used by request-batch admission.
+    # This is populated by a pipeline preprocessor before the request reaches
+    # the scheduler; ``None`` keeps the default behavior for other pipelines.
+    batch_compatibility_key: tuple[Any, ...] | None = None
 
     def __post_init__(self):
         """Initialize dependent fields after dataclass initialization."""

vllm_omni/diffusion/sched/interface.py

@@ -88,9 +88,10 @@ class RequestBatchSamplingParamsKey:
     """Request level Batch-compatibility key derived from ``OmniDiffusionSamplingParams``.
 
     Only request-batch-wide fields belong here. Request-local values such as
-    seeds, generators, latent tensors, timesteps, and pipeline-specific
+    seeds, generators, latent tensors, timesteps, and request-local
     ``extra_args`` are read per request from
-    ``DiffusionRequestBatch.sampling_params_list``.
+    ``DiffusionRequestBatch.sampling_params_list``. Structural ``extra_args``
+    that affect batching are listed explicitly below.
     """
 
     # Spatial / temporal shape.
@@ -131,6 +132,16 @@ class RequestBatchSamplingParamsKey:
     # the model and must remain distinct from explicit ``lossless``.
     quality: str | None = None
 
+    # Wan scheduler structure is carried through extra_args. Requests using
+    # different solvers or flow shifts must not share a request batch.
+    sample_solver: str | None = None
+    flow_shift: float | None = None
+
+    # Pipeline-specific condition structure populated during preprocessing.
+    # It prevents independently valid requests with incompatible conditions
+    # from being admitted to the same request batch.
+    condition_key: tuple[Any, ...] | None = None
+
     # LoRA identity.
     lora_int_id: int | None = None
     lora_scale: float = 1.0

vllm_omni/diffusion/sched/request_scheduler.py

@@ -22,7 +22,21 @@ if TYPE_CHECKING:
 # on sampling params, so it must be resolved separately from the bulk lookup.
 _REQUEST_BATCH_SAMPLING_PARAMS_KEY_FIELD_NAMES = frozenset(
     field.name for field in fields(RequestBatchSamplingParamsKey)
-) - {"lora_int_id"}
+) - {"condition_key", "flow_shift", "lora_int_id", "sample_solver"}
+
+
+def _normalize_explicit_sample_solver(value: object | None) -> str | None:
+    """Normalize an explicitly provided solver without selecting a default."""
+    if value is None:
+        return None
+    return str(value).strip().lower()
+
+
+def _normalize_explicit_flow_shift(value: object | None) -> float | None:
+    """Normalize an explicitly provided flow shift without selecting a default."""
+    if value is None:
+        return None
+    return float(value)
 
 
 def build_request_batch_sampling_params_key(request: OmniDiffusionRequest) -> RequestBatchSamplingParamsKey:
@@ -31,6 +45,13 @@ def build_request_batch_sampling_params_key(request: OmniDiffusionRequest) -> Re
     # LoRA identity is optional on sampling params (and on test stubs).
     lora_request = getattr(sampling, "lora_request", None)
     key_kwargs = {name: getattr(sampling, name) for name in _REQUEST_BATCH_SAMPLING_PARAMS_KEY_FIELD_NAMES}
+    extra_args = sampling.extra_args or {}
+    # Match pipeline resolution for explicit overrides, but preserve None:
+    # pipeline/engine defaults are configuration-dependent and must not be
+    # inferred while building the request-batch key.
+    key_kwargs["sample_solver"] = _normalize_explicit_sample_solver(extra_args.get("sample_solver"))
+    key_kwargs["flow_shift"] = _normalize_explicit_flow_shift(extra_args.get("flow_shift"))
+    key_kwargs["condition_key"] = getattr(request, "batch_compatibility_key", None)
     key_kwargs["lora_int_id"] = lora_request.lora_int_id if lora_request is not None else None
     return RequestBatchSamplingParamsKey(**key_kwargs)
 

vllm_omni/engine/stage_engine_startup.py

@@ -1558,6 +1558,7 @@ def launch_diffusion_stage_replica(
     from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
 
     od_config = build_diffusion_config(model, stage_config, metadata)
+    od_config.max_num_seqs = batch_size
     parallel_config = getattr(od_config, "parallel_config", None)
     world_size = getattr(parallel_config, "world_size", 1)
     try:

vllm_omni/engine/stage_init_utils.py

@@ -1441,6 +1441,7 @@ def initialize_diffusion_stage(
     from vllm_omni.diffusion.stage_diffusion_client import create_diffusion_client
 
     od_config = build_diffusion_config(model, stage_cfg, metadata)
+    od_config.max_num_seqs = batch_size
     return create_diffusion_client(model, od_config, metadata, stage_init_timeout, batch_size, use_inline)