[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>

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:
-        prompt: str | None = None
-        negative_prompt: str | None = None
-        prompt_embeds: torch.Tensor | None = None
-        negative_prompt_embeds: torch.Tensor | None = None
-        if len(req.prompts) > 1:
-            raise ValueError(
-                """This model only supports a single prompt, not a batched request.""",
-                """Please pass in a single prompt object or string, or a single-item list.""",
-            )
-        if len(req.prompts) == 1:
-            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")
+    def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]:
+        sampling_params_list = req.sampling_params_list
+        common = sampling_params_list[0]
+        prompt_texts = [prompt if isinstance(prompt, str) else (prompt.get("prompt") or "") for prompt in req.prompts]
+        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"]
+        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:
-            raise ValueError("Prompt is required for Wan2.2 generation.")
+        if prompt is not None and not all(prompt):
+            raise ValueError("Prompt is required for Wan2.2 generation when prompt_embeds are not provided.")
 
-        height = req.sampling_params.height or 480
-        width = req.sampling_params.width or 832
-        num_frames = req.sampling_params.num_frames or 81
+        height = common.height or 480
+        width = common.width or 832
+        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.sampling_params.generator
-        if generator is None and req.sampling_params.seed is not None:
-            generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed)
+        generator = req.collate_request_generators(num_outputs_per_prompt, None)
+        request_latents = req.collate_request_tensors("latents", None)
 
         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(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0,
-            num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1,
-            max_sequence_length=req.sampling_params.max_sequence_length or 512,
-            device=device,
-            dtype=dtype,
-        )
+        do_classifier_free_guidance = guidance_low > 1.0 or guidance_high > 1.0
+        if prompt_embeds is None:
+            prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                do_classifier_free_guidance=do_classifier_free_guidance,
+                num_videos_per_prompt=num_outputs_per_prompt,
+                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)
-        flow_shift = resolve_wan_flow_shift(req, self.od_config)
+        first_request = req.requests[0]
+        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
-        raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
-        if isinstance(raw_image, list):
-            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 raw_image is None:
-            image = None
-        elif isinstance(raw_image, str):
-            image = PIL.Image.open(raw_image)
-        else:
-            image = cast(PIL.Image.Image | torch.Tensor, raw_image)
+        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 isinstance(raw_image, list):
+                if len(raw_image) > 1:
+                    logger.warning("Received multiple images for one Wan request; using only the first image.")
+                raw_image = raw_image[0] if raw_image else None
+            if isinstance(raw_image, str):
+                raw_image = PIL.Image.open(raw_image)
+            images.append(cast(PIL.Image.Image | torch.Tensor | None, 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
-            if isinstance(image, PIL.Image.Image):
-                image = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
-                image_tensor = video_processor.preprocess(image, height=height, width=width)
-            else:
-                image_tensor = image
+            image_tensors = []
+            for image in images:
+                assert image is not None
+                if isinstance(image, PIL.Image.Image):
+                    image = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
+                    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(
-            output=output, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
+        return split_diffusion_output_by_request(
+            DiffusionOutput(
+                output=output,
+                stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
+            ),
+            req,
+            num_outputs_per_prompt=num_outputs_per_prompt,
         )
 
     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:
-        prompt: str | None = None
-        negative_prompt: str | None = None
-        prompt_embeds: torch.Tensor | None = None
-        negative_prompt_embeds: torch.Tensor | None = None
+    def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]:
+        sampling_params_list = req.sampling_params_list
+        common = sampling_params_list[0]
+        prompt_texts = [prompt if isinstance(prompt, str) else (prompt.get("prompt") or "") for prompt in req.prompts]
+        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:
-            raise ValueError(
-                """This model only supports a single prompt, not a batched request.""",
-                """Please pass in a single prompt object or string, or a single-item list.""",
-            )
-        if len(req.prompts) == 1:
-            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")
+        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 prompt is not None and not all(prompt):
+            raise ValueError("Prompt is required for Wan2.2 I2V generation when prompt_embeds are not provided.")
 
-        if not prompt:
-            raise ValueError("Prompt is required for Wan2.2 generation.")
+        images: list[PIL.Image.Image | torch.Tensor] = []
+        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
-        multi_modal_data = req.prompts[0].get("multi_modal_data", {}) if not isinstance(req.prompts[0], str) else None
-        raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
-        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 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 = multi_modal_data.get("last_image")
+            if isinstance(last_image, str):
+                last_image = PIL.Image.open(last_image).convert("RGB")
+            last_images.append(cast(PIL.Image.Image | torch.Tensor | None, last_image))
+        if any(image is not None for image in last_images) and not all(image is not None for image in last_images):
+            raise ValueError("Cannot batch I2V requests with a mix of provided and missing last_image conditions.")
 
-        last_image: PIL.Image.Image | torch.Tensor | None = None
-        if multi_modal_data is not None:
-            last_image = multi_modal_data.get("last_image", None)
+        height = common.height or 480
+        width = common.width or 832
+        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
-        width = req.sampling_params.width or 832
-        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"
+        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=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.sampling_params.generator
-        if generator is None and req.sampling_params.seed is not None:
-            generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed)
+        generator = req.collate_request_generators(num_outputs_per_prompt, None)
+        request_latents = req.collate_request_tensors("latents", None)
 
         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(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0,
-            num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1,
-            max_sequence_length=req.sampling_params.max_sequence_length or 512,
-            device=device,
-            dtype=dtype,
-        )
+        do_classifier_free_guidance = guidance_low > 1.0 or guidance_high > 1.0
+        if prompt_embeds is None:
+            prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                do_classifier_free_guidance=do_classifier_free_guidance,
+                num_videos_per_prompt=num_outputs_per_prompt,
+                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 last_image is None:
-                    image_embeds = self.encode_image(image, device)
-                else:
-                    image_embeds = self.encode_image([image, last_image], device)
-            image_embeds = image_embeds.repeat(batch_size, 1, 1)
+            if all(last_image is None for last_image in last_images):
+                image_embeds = self.encode_image(images, device)
+                image_embeds = image_embeds.repeat_interleave(num_outputs_per_prompt, dim=0)
+            else:
+                image_pairs = []
+                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)
-        flow_shift = resolve_wan_flow_shift(req, self.od_config)
+        first_request = req.requests[0]
+        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 = TF.to_tensor(image).to(device)
-            image_tensor = video_processor.preprocess(image, height=height, width=width)
-        else:
-            image_tensor = image
-        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)
+        image_tensors = []
+        for image in images:
+            if isinstance(image, PIL.Image.Image):
+                image = TF.to_tensor(image).to(device)
+                image_tensor = video_processor.preprocess(image, height=height, width=width)
             else:
-                last_image_tensor = last_image
-            last_image_tensor = last_image_tensor.to(device=device, dtype=torch.float32)
-        else:
-            last_image_tensor = None
+                image_tensor = image.unsqueeze(0) if image.ndim == 3 else image
+            image_tensors.append(image_tensor.to(device=device, dtype=torch.float32))
+        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)
+
+        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(
-            output=output, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
+        return split_diffusion_output_by_request(
+            DiffusionOutput(
+                output=output,
+                stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
+            ),
+            req,
+            num_outputs_per_prompt=num_outputs_per_prompt,
         )
 
     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.sample.squeeze(0)
+            return result[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,
-        negative_audio_emb: torch.Tensor | None,
+        positive_audio_emb: dict[str, torch.Tensor],
+        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)]
-                timestep = t.unsqueeze(0).to(device) if t.dim() == 0 else t.to(device)
+                latent_model_input = latents.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,
-                )[0].squeeze(0)
+                    generator=clip_generator if isinstance(clip_generator, torch.Generator) else None,
+                )[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 ----
-        if len(req.prompts) > 1:
-            raise ValueError("S2V only supports a single prompt per request.")
-
-        if len(req.prompts) == 1:
-            prompt_data = req.prompts[0]
-            prompt = prompt_data if isinstance(prompt_data, str) else prompt_data.get("prompt")
-            negative_prompt = None if isinstance(prompt_data, str) else prompt_data.get("negative_prompt")
+        sampling_params_list = req.sampling_params_list
+        common = sampling_params_list[0]
+        prompt_fields = DiffusionRequestBatch.collate_prompt_field_map(
+            req.prompts,
+            {
+                "prompt_embeds": prompt_embeds,
+                "negative_prompt_embeds": negative_prompt_embeds,
+            },
+        )
+        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:
-                raw_image = multi_modal_data.get("image", None)
-                if isinstance(raw_image, str):
-                    image = PIL.Image.open(raw_image).convert("RGB")
-                elif isinstance(raw_image, PIL.Image.Image):
-                    image = raw_image
+            prompt_texts.append(
+                prompt_data if isinstance(prompt_data, str) else (prompt_data.get("prompt") or prompt or "")
+            )
+            request_negative_prompt = (
+                negative_prompt if isinstance(prompt_data, str) else prompt_data.get("negative_prompt", negative_prompt)
+            )
+            negative_prompts.append(request_negative_prompt or _S2V_DEFAULT_NEG_PROMPT)
 
-            if audio_path is None:
-                audio_path = additional_info.get("audio_path")
-                if audio_path is None:
-                    audio_path = multi_modal_data.get("audio")
+            raw_image = multi_modal_data.get("image", image)
+            if isinstance(raw_image, list):
+                raw_image = raw_image[0] if raw_image else None
+            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:
-                pose_video = additional_info.get("pose_video")
+            request_audio_path = additional_info.get("audio_path", audio_path)
+            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)
-
-        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
+        height = common.height or height
+        width = common.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:
-            guidance_scale = req.sampling_params.guidance_scale
+        if common.guidance_scale_provided:
+            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
-
-        if generator is None:
-            generator = req.sampling_params.generator
-        seed = req.sampling_params.seed
-        if generator is None and seed is not None:
-            generator = torch.Generator(device=device).manual_seed(seed)
-        if seed is None:
-            seed = 0
+        request_generators = req.collate_request_generators(num_outputs_per_prompt, None)
+        batch_size = req.num_reqs * num_outputs_per_prompt
+        if isinstance(request_generators, list):
+            generators = request_generators
+        else:
+            generators = [request_generators] * batch_size
+        seeds = [sampling.seed for sampling in sampling_params_list for _ in range(num_outputs_per_prompt)]
+        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,
-                negative_prompt=negative_prompt,
+                prompt=prompt_texts,
+                negative_prompt=negative_prompts,
                 do_classifier_free_guidance=self.do_classifier_free_guidance,
-                num_videos_per_prompt=1,
-                max_sequence_length=512,
+                num_videos_per_prompt=num_outputs_per_prompt,
+                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_path, infer_frames=infer_frames, device=device, dtype=dtype
-        )
-        if num_repeat is None or num_repeat > audio_num_repeat:
-            num_repeat = audio_num_repeat
+        audio_embeddings: list[torch.Tensor] = []
+        effective_repeats: list[int] = []
+        target_video_frames: list[int] = []
+        for request_audio_path, requested_repeat in zip(audio_paths, requested_repeats):
+            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_waveform = audio_path.astype(np.float32)
-            raw_audio_sr = 16000
-        else:
-            raw_audio_waveform, raw_audio_sr = load_audio(audio_path, sr=None, mono=True)
+        raw_audio_waveforms: list[np.ndarray] = []
+        raw_audio_sample_rates: list[int] = []
+        for request_audio_path in audio_paths:
+            assert request_audio_path is not None
+            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]
-            motion_pixels[:, :, -6:] = ref_tensor.unsqueeze(0).unsqueeze(2).expand(-1, -1, 6, -1, -1)
+            ref_tensors = torch.stack(
+                [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_seed = seed + r
-            clip_generator = torch.Generator(device=device).manual_seed(clip_seed)
-
+            clip_generators = _make_clip_generators(seeds, generators, r, device)
             # -- Noise --
-            latents = self.prepare_latents(
-                infer_frames=infer_frames,
-                height=height,
-                width=width,
-                motion_frames=motion_frames,
-                lat_motion_frames=lat_motion_frames,
-                dtype=dtype,
-                device=device,
-                generator=clip_generator,
-            )
+            if r == 0 and request_latents is not None:
+                latents = request_latents.to(device=device, dtype=dtype)
+            else:
+                latents = torch.stack(
+                    [
+                        self.prepare_latents(
+                            infer_frames=infer_frames,
+                            height=height,
+                            width=width,
+                            motion_frames=motion_frames,
+                            lat_motion_frames=lat_motion_frames,
+                            dtype=dtype,
+                            device=device,
+                            generator=clip_generator,
+                        )
+                        for clip_generator in clip_generators
+                    ]
+                )
 
             # -- Scheduler --
             self.scheduler.set_timesteps(num_steps, device=device)
@@ -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(
-            output=(output, raw_audio_waveform, raw_audio_sr),
-            stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
+        outputs = split_diffusion_output_by_request(
+            DiffusionOutput(
+                output=output,
+                stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
+            ),
+            req,
+            num_outputs_per_prompt=num_outputs_per_prompt,
         )
+        for request_output, raw_audio_waveform in zip(outputs, raw_audio_waveforms):
+            request_output.output = (request_output.output, raw_audio_waveform, raw_audio_sr)
+        return outputs
 
     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
-        if len(req.prompts) > 1:
-            raise ValueError(
-                "This model only supports a single prompt, not a batched request. "
-                "Please pass in a single prompt object or string, or a single-item list."
-            )
-
-        reference_images = None
-        source_video = None
-        source_mask = None
-
-        if len(req.prompts) == 1:
-            first_prompt = req.prompts[0]
-            if isinstance(first_prompt, str):
-                prompt = first_prompt
-            else:
-                prompt = first_prompt.get("prompt")
-                negative_prompt = negative_prompt or first_prompt.get("negative_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:
+        sampling_params_list = req.sampling_params_list
+        common = sampling_params_list[0]
+        prompt_texts = [item if isinstance(item, str) else (item.get("prompt") or "") for item in req.prompts]
+        negative_prompts = [None if isinstance(item, str) else item.get("negative_prompt") for item in req.prompts]
+        prompt_fields = DiffusionRequestBatch.collate_prompt_field_map(
+            req.prompts,
+            {
+                "prompt_embeds": prompt_embeds,
+                "negative_prompt_embeds": negative_prompt_embeds,
+            },
+        )
+        prompt_embeds = prompt_fields["prompt_embeds"]
+        negative_prompt_embeds = prompt_fields["negative_prompt_embeds"]
+        prompt = prompt_texts if prompt_embeds is None else None
+        negative_prompt = 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 prompt is not None and not all(prompt):
             raise ValueError("Prompt or prompt_embeds is required for VACE generation.")
 
-        height = req.sampling_params.height or height
-        width = req.sampling_params.width or width
-        num_frames = req.sampling_params.num_frames or frame_num
-        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
-        generator = req.sampling_params.generator or generator
+        reference_images_list = []
+        source_videos = []
+        source_masks = []
+        for request_prompt in req.prompts:
+            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:
-            guidance_scale = req.sampling_params.guidance_scale
+        height = common.height or height
+        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(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            height=height,
-            width=width,
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            video=source_video,
-            mask=source_mask,
-            reference_images=reference_images,
-        )
+        for index in range(req.num_reqs):
+            self.check_inputs(
+                prompt=None if prompt is None else prompt[index],
+                negative_prompt=None if negative_prompt is None else negative_prompt[index],
+                height=height,
+                width=width,
+                prompt_embeds=None if prompt_embeds is None else prompt_embeds[index],
+                negative_prompt_embeds=None if negative_prompt_embeds is None else negative_prompt_embeds[index],
+                video=source_videos[index],
+                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,
-                max_sequence_length=req.sampling_params.max_sequence_length or 512,
+                num_videos_per_prompt=num_outputs_per_prompt,
+                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="",
-                    negative_prompt=None,
+                    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,
                 )
@@ -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(
-                video=source_video,
-                mask=source_mask,
-                reference_images=reference_images,
-                height=height,
-                width=width,
-                num_frames=num_frames,
-                dtype=dtype,
-                device=device,
-            )
+            contexts = []
+            reference_counts = []
+            for index, (source_video, source_mask, reference_images) in enumerate(
+                zip(source_videos, source_masks, reference_images_list)
+            ):
+                video, mask, ref_images_processed = self.preprocess_conditions(
+                    video=source_video,
+                    mask=source_mask,
+                    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)
-            mask_encoded = self.prepare_masks(mask, ref_images_processed)
-
-            # Unified VACE context: [video_latents, mask] along channels -> [B, C, T, H, W]
-            vace_context = torch.cat([conditioning_latents, mask_encoded], dim=1)
-
-            num_reference_images = len(ref_images_processed[0]) if ref_images_processed else 0
+            if len(set(reference_counts)) != 1:
+                raise ValueError(
+                    f"Batched VACE conditions must have the same number of reference images; got {reference_counts}."
+                )
+            vace_context = DiffusionRequestBatch.collate_tensors(contexts, "VACE condition", None)
+            assert vace_context is not None
+            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 = (
-                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_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."""