[diffusion] fix: map each prompt to corresponding image in multi-prompt scenario (#20081)

Signed-off-by: Lancer <maruixiang6688@gmail.com>
Co-authored-by: Mick <mickjagger19@icloud.com>

python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py

@@ -56,6 +56,37 @@ def qwen_image_postprocess_text(outputs, _text_inputs, drop_idx=34):
     return prompt_embeds
 
 
+def _normalize_prompt_list(prompt):
+    return [prompt] if isinstance(prompt, str) else prompt
+
+
+def _normalize_image_list(images):
+    if images is None:
+        return []
+    return images if isinstance(images, list) else [images]
+
+
+def _build_qwen_edit_image_prompt(num_images: int) -> str:
+    img_prompt_template = "Picture {}: <|vision_start|><|image_pad|><|vision_end|>"
+    return "".join(img_prompt_template.format(i + 1) for i in range(num_images))
+
+
+def _resolve_qwen_edit_per_prompt_images(prompt_list, image_list):
+    if len(prompt_list) <= 1:
+        return [image_list]
+
+    if len(image_list) <= 1:
+        return [list(image_list) for _ in prompt_list]
+
+    if len(image_list) != len(prompt_list):
+        raise ValueError(
+            "QwenImageEditPlus expects either one shared condition image or "
+            "the same number of condition images and prompts."
+        )
+
+    return [[image] for image in image_list]
+
+
 # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._pack_latents
 def _pack_latents(latents, batch_size, num_channels_latents, height, width):
     latents = latents.view(
@@ -372,8 +403,14 @@ class QwenImageEditPlusPipelineConfig(QwenImageEditPipelineConfig):
 
     def prepare_image_processor_kwargs(self, batch, neg=False) -> dict:
         prompt = batch.prompt if not neg else batch.negative_prompt
-        prompt_list = [prompt] if isinstance(prompt, str) else prompt
+        if not prompt:
-        image_list = batch.condition_image
+            return {}
+
+        prompt_list = _normalize_prompt_list(prompt)
+        image_list = _normalize_image_list(batch.condition_image)
+        per_prompt_images = _resolve_qwen_edit_per_prompt_images(
+            prompt_list, image_list
+        )
 
         prompt_template_encode = (
             "<|im_start|>system\nDescribe the key features of the input image "
@@ -384,13 +421,14 @@ class QwenImageEditPlusPipelineConfig(QwenImageEditPipelineConfig):
             "<|im_start|>user\n{}<|im_end|>\n"
             "<|im_start|>assistant\n"
         )
-        img_prompt_template = "Picture {}: <|vision_start|><|image_pad|><|vision_end|>"
+        txt = [
-        if isinstance(image_list, list):
+            prompt_template_encode.format(
-            base_img_prompt = ""
+                _build_qwen_edit_image_prompt(len(prompt_images)) + prompt_text
-            for i, img in enumerate(image_list):
+            )
-                base_img_prompt += img_prompt_template.format(i + 1)
+            for prompt_text, prompt_images in zip(prompt_list, per_prompt_images)
-        txt = [prompt_template_encode.format(base_img_prompt + p) for p in prompt_list]
+        ]
-        return dict(text=txt, padding=True)
+
+        return dict(text=txt, padding=True, per_prompt_images=per_prompt_images)
 
     def prepare_calculated_size(self, image):
         return self.calculate_vae_image_size(image, image.width, image.height)
@@ -510,7 +548,6 @@ class QwenImageLayeredPipelineConfig(QwenImageEditPipelineConfig):
         assert batch_size == 1
         height = batch.height
         width = batch.width
-        image_size = batch.original_condition_image_size
 
         vae_scale_factor = self.get_vae_scale_factor()
 

python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py

@@ -155,6 +155,24 @@ class DiffGenerator:
             f"{self.server_args.scheduler_endpoint}."
         )
 
+    @staticmethod
+    def _resolve_image_paths_per_prompt(
+        prompts: list[str], image_paths: str | list[str] | None
+    ) -> list[str | list[str] | None]:
+        if len(prompts) <= 1:
+            return [image_paths]
+
+        if not isinstance(image_paths, list) or len(image_paths) <= 1:
+            return [image_paths for _ in prompts]
+
+        if len(image_paths) != len(prompts):
+            raise ValueError(
+                "When using multiple prompts with multiple input images, "
+                "provide either one shared image or exactly one image per prompt."
+            )
+
+        return [[image_path] for image_path in image_paths]
+
     def generate(
         self,
         sampling_params_kwargs: dict | None = None,
@@ -181,11 +199,16 @@ class DiffGenerator:
         )
 
         requests: list[Req] = []
-        for p in prompts:
+        image_paths_per_prompt = self._resolve_image_paths_per_prompt(
+            prompts, sampling_params_orig.image_path
+        )
+
+        for i, p in enumerate(prompts):
             sampling_params = dataclasses.replace(
                 sampling_params_orig,
                 prompt=p,
                 output_file_name=user_output_file_name,
+                image_path=image_paths_per_prompt[i],
             )
             sampling_params._set_output_file_name()
             req = prepare_request(

python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py

@@ -105,15 +105,32 @@ class ImageEncodingStage(PipelineStage):
         cuda_device = get_local_torch_device()
 
         self.load_model()
-        image = batch.condition_image
 
         image_processor_kwargs = (
             server_args.pipeline_config.prepare_image_processor_kwargs(batch)
         )
+        per_prompt_images = image_processor_kwargs.pop("per_prompt_images", None)
+        texts = image_processor_kwargs.pop("text", None)
+
+        if per_prompt_images is None:
+            per_prompt_images = [batch.condition_image]
+            texts = [None] if texts is None else texts
+
+        all_prompt_embeds = []
+        all_neg_prompt_embeds = []
+
+        for idx, prompt_images in enumerate(per_prompt_images):
+            if not prompt_images:
+                continue
+
+            cur_kwargs = image_processor_kwargs.copy()
+            if texts and idx < len(texts):
+                cur_kwargs["text"] = [texts[idx]]
 
             image_inputs = self.image_processor(
-            images=image, return_tensors="pt", **image_processor_kwargs
+                images=prompt_images, return_tensors="pt", **cur_kwargs
             ).to(cuda_device)
+
             if self.image_encoder:
                 # if an image encoder is provided
                 with set_forward_context(current_timestep=0, attn_metadata=None):
@@ -121,8 +138,9 @@ class ImageEncodingStage(PipelineStage):
                         **image_inputs,
                         **server_args.pipeline_config.image_encoder_extra_args,
                     )
-                image_embeds = server_args.pipeline_config.postprocess_image(outputs)
+                    image_embeds = server_args.pipeline_config.postprocess_image(
-
+                        outputs
+                    )
                 batch.image_embeds.append(image_embeds)
             elif self.text_encoder:
                 # if a text encoder is provided, e.g. Qwen-Image-Edit
@@ -133,9 +151,14 @@ class ImageEncodingStage(PipelineStage):
                             batch, neg=True
                         )
                     )
-
+                    neg_image_processor_kwargs.pop("per_prompt_images", None)
+                    neg_texts = neg_image_processor_kwargs.pop("text", None)
+                    if neg_texts and idx < len(neg_texts):
+                        neg_image_processor_kwargs["text"] = [neg_texts[idx]]
                     neg_image_inputs = self.image_processor(
-                    images=image, return_tensors="pt", **neg_image_processor_kwargs
+                        images=prompt_images,
+                        return_tensors="pt",
+                        **neg_image_processor_kwargs,
                     ).to(cuda_device)
 
                 with set_forward_context(current_timestep=0, attn_metadata=None):
@@ -154,15 +177,20 @@ class ImageEncodingStage(PipelineStage):
                             image_grid_thw=neg_image_inputs.image_grid_thw,
                             output_hidden_states=True,
                         )
-            batch.prompt_embeds.append(
+
+                all_prompt_embeds.append(
                     self.encoding_qwen_image_edit(outputs, image_inputs)
                 )
-
                 if batch.do_classifier_free_guidance:
-                batch.negative_prompt_embeds.append(
+                    all_neg_prompt_embeds.append(
                         self.encoding_qwen_image_edit(neg_outputs, neg_image_inputs)
                     )
 
+        if all_prompt_embeds:
+            batch.prompt_embeds.append(torch.cat(all_prompt_embeds, dim=0))
+        if all_neg_prompt_embeds:
+            batch.negative_prompt_embeds.append(torch.cat(all_neg_prompt_embeds, dim=0))
+
         self.offload_model()
 
         return batch

python/sglang/multimodal_gen/test/cli/test_generate_common.py

@@ -58,6 +58,24 @@ class CLIBase(unittest.TestCase):
     height: int = 720
     output_path: str = "test_outputs"
 
+    def setUp(self):
+        super().setUp()
+        if not os.path.exists(self.output_path):
+            os.makedirs(self.output_path, exist_ok=True)
+        if os.path.exists(self.output_path):
+            for f in os.listdir(self.output_path):
+                path = os.path.join(self.output_path, f)
+                if os.path.isfile(path):
+                    os.remove(path)
+
+    def tearDown(self):
+        super().tearDown()
+        if os.path.exists(self.output_path):
+            for f in os.listdir(self.output_path):
+                path = os.path.join(self.output_path, f)
+                if os.path.isfile(path):
+                    os.remove(path)
+
     def get_base_command(self):
         return [
             "sglang",

python/sglang/multimodal_gen/test/cli/test_generate_i2i.py

@@ -0,0 +1,132 @@
+import os
+import unittest
+
+from PIL import Image
+
+from sglang.multimodal_gen.configs.sample.sampling_params import DataType
+from sglang.multimodal_gen.test.cli.test_generate_common import CLIBase, run_command
+from sglang.multimodal_gen.test.test_utils import (
+    DEFAULT_QWEN_IMAGE_EDIT_2511_MODEL_NAME_FOR_TEST,
+    check_image_size,
+)
+
+
+class TestQwenImageEditI2I(CLIBase):
+    model_path: str = DEFAULT_QWEN_IMAGE_EDIT_2511_MODEL_NAME_FOR_TEST
+    data_type: DataType = DataType.IMAGE
+    width: int = 512
+    height: int = 512
+
+    test_image_urls = [
+        "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/edit2509/edit2509_1.jpg",
+        "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/edit2509/edit2509_2.jpg",
+    ]
+
+    def get_base_command(self):
+        return [
+            "sglang",
+            "generate",
+            "--save-output",
+            "--log-level=info",
+            f"--width={self.width}",
+            f"--height={self.height}",
+            f"--output-path={self.output_path}",
+        ]
+
+    def verify_multi_output(self, name: str, num_outputs: int):
+        output_files = []
+        try:
+            all_files = os.listdir(self.output_path)
+            ext = self.data_type.get_default_extension()
+            for f in all_files:
+                if f.endswith(f".{ext}"):
+                    output_files.append(f)
+
+            self.assertEqual(
+                len(output_files),
+                num_outputs,
+                f"Expected {num_outputs} output files, found {len(output_files)}: {output_files}",
+            )
+
+            for f in output_files:
+                path = os.path.join(self.output_path, f)
+                with Image.open(path) as image:
+                    check_image_size(self, image, self.width, self.height)
+        finally:
+            for f in output_files:
+                path = os.path.join(self.output_path, f)
+                if os.path.exists(path):
+                    os.remove(path)
+
+    def test_single_prompt_single_image(self):
+        """Case 1: Single prompt + single image."""
+        name = "single_prompt_single_image"
+
+        command = self.get_base_command() + [
+            f"--model-path={self.model_path}",
+            "--prompt",
+            "Add a red hat",
+            "--image-path",
+            self.test_image_urls[0],
+        ]
+
+        succeed = run_command(command)
+        self.assertTrue(succeed, f"{name} command failed")
+        self.verify_multi_output(name, 1)
+
+    def test_single_prompt_multi_image(self):
+        """Case 2: Single prompt + multiple images (image composition)."""
+        name = "single_prompt_multi_image"
+
+        command = self.get_base_command() + [
+            f"--model-path={self.model_path}",
+            "--prompt",
+            "Combine both images",
+            "--image-path",
+            *self.test_image_urls,
+        ]
+
+        succeed = run_command(command)
+        self.assertTrue(succeed, f"{name} command failed")
+        self.verify_multi_output(name, 1)
+
+    def test_multi_prompt_multi_image(self):
+        """Case 3: Multiple prompts + multiple images (image editing)."""
+        name = "multi_prompt_multi_image"
+
+        command = self.get_base_command() + [
+            f"--model-path={self.model_path}",
+            "--prompt",
+            "Convert to oil painting style",
+            "Convert to watercolor style",
+            "--image-path",
+            *self.test_image_urls,
+        ]
+
+        succeed = run_command(command)
+        self.assertTrue(succeed, f"{name} command failed")
+        self.verify_multi_output(name, 2)
+
+    def test_multi_prompt_single_image(self):
+        """Case 4: Multiple prompts + single image (image editing)."""
+        name = "multi_prompt_single_image"
+
+        command = self.get_base_command() + [
+            f"--model-path={self.model_path}",
+            "--prompt",
+            "Add a red hat",
+            "Change to blue background",
+            "--image-path",
+            self.test_image_urls[0],
+        ]
+
+        succeed = run_command(command)
+        self.assertTrue(succeed, f"{name} command failed")
+        self.verify_multi_output(name, 2)
+
+
+del CLIBase
+
+
+if __name__ == "__main__":
+    unittest.main()