From 8cd1de3354d032bb1527c0db7b2a944e838fc368 Mon Sep 17 00:00:00 2001 From: Lancer Date: Tue, 10 Mar 2026 16:58:21 +0800 Subject: [PATCH] [diffusion] fix: map each prompt to corresponding image in multi-prompt scenario (#20081) Signed-off-by: Lancer Co-authored-by: Mick --- .../configs/pipeline_configs/qwen_image.py | 57 ++++++-- .../entrypoints/diffusion_generator.py | 25 +++- .../pipelines_core/stages/image_encoding.py | 118 ++++++++++------ .../test/cli/test_generate_common.py | 18 +++ .../test/cli/test_generate_i2i.py | 132 ++++++++++++++++++ 5 files changed, 294 insertions(+), 56 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/cli/test_generate_i2i.py diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py b/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py index 08d20c958..991fb6ea6 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py +++ b/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 - image_list = batch.condition_image + if not prompt: + 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|>" - if isinstance(image_list, list): - base_img_prompt = "" - for i, img in enumerate(image_list): - base_img_prompt += img_prompt_template.format(i + 1) - txt = [prompt_template_encode.format(base_img_prompt + p) for p in prompt_list] - return dict(text=txt, padding=True) + txt = [ + prompt_template_encode.format( + _build_qwen_edit_image_prompt(len(prompt_images)) + prompt_text + ) + for prompt_text, prompt_images in zip(prompt_list, per_prompt_images) + ] + + 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() diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index b37e1f8bb..2f71ad1d1 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/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( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py index 988cad561..861bdda97 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py @@ -105,63 +105,91 @@ 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) - image_inputs = self.image_processor( - images=image, return_tensors="pt", **image_processor_kwargs - ).to(cuda_device) - if self.image_encoder: - # if an image encoder is provided - with set_forward_context(current_timestep=0, attn_metadata=None): - outputs = self.image_encoder( - **image_inputs, - **server_args.pipeline_config.image_encoder_extra_args, - ) - image_embeds = server_args.pipeline_config.postprocess_image(outputs) + if per_prompt_images is None: + per_prompt_images = [batch.condition_image] + texts = [None] if texts is None else texts - batch.image_embeds.append(image_embeds) - elif self.text_encoder: - # if a text encoder is provided, e.g. Qwen-Image-Edit - # 1. neg prompt embeds - if batch.do_classifier_free_guidance: - neg_image_processor_kwargs = ( - server_args.pipeline_config.prepare_image_processor_kwargs( - batch, neg=True + 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=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): + outputs = self.image_encoder( + **image_inputs, + **server_args.pipeline_config.image_encoder_extra_args, ) - ) - - neg_image_inputs = self.image_processor( - images=image, return_tensors="pt", **neg_image_processor_kwargs - ).to(cuda_device) - - with set_forward_context(current_timestep=0, attn_metadata=None): - outputs = self.text_encoder( - input_ids=image_inputs.input_ids, - attention_mask=image_inputs.attention_mask, - pixel_values=image_inputs.pixel_values, - image_grid_thw=image_inputs.image_grid_thw, - output_hidden_states=True, - ) + 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 + # 1. neg prompt embeds if batch.do_classifier_free_guidance: - neg_outputs = self.text_encoder( - input_ids=neg_image_inputs.input_ids, - attention_mask=neg_image_inputs.attention_mask, - pixel_values=neg_image_inputs.pixel_values, - image_grid_thw=neg_image_inputs.image_grid_thw, + neg_image_processor_kwargs = ( + server_args.pipeline_config.prepare_image_processor_kwargs( + 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=prompt_images, + return_tensors="pt", + **neg_image_processor_kwargs, + ).to(cuda_device) + + with set_forward_context(current_timestep=0, attn_metadata=None): + outputs = self.text_encoder( + input_ids=image_inputs.input_ids, + attention_mask=image_inputs.attention_mask, + pixel_values=image_inputs.pixel_values, + image_grid_thw=image_inputs.image_grid_thw, output_hidden_states=True, ) - batch.prompt_embeds.append( - self.encoding_qwen_image_edit(outputs, image_inputs) - ) + if batch.do_classifier_free_guidance: + neg_outputs = self.text_encoder( + input_ids=neg_image_inputs.input_ids, + attention_mask=neg_image_inputs.attention_mask, + pixel_values=neg_image_inputs.pixel_values, + image_grid_thw=neg_image_inputs.image_grid_thw, + output_hidden_states=True, + ) - if batch.do_classifier_free_guidance: - batch.negative_prompt_embeds.append( - self.encoding_qwen_image_edit(neg_outputs, neg_image_inputs) + all_prompt_embeds.append( + self.encoding_qwen_image_edit(outputs, image_inputs) ) + if batch.do_classifier_free_guidance: + 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() diff --git a/python/sglang/multimodal_gen/test/cli/test_generate_common.py b/python/sglang/multimodal_gen/test/cli/test_generate_common.py index 4bcbf8ab1..94d97ff3f 100644 --- a/python/sglang/multimodal_gen/test/cli/test_generate_common.py +++ b/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", diff --git a/python/sglang/multimodal_gen/test/cli/test_generate_i2i.py b/python/sglang/multimodal_gen/test/cli/test_generate_i2i.py new file mode 100644 index 000000000..c7c4cbbfb --- /dev/null +++ b/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()