chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Lancer <maruixiang6688@gmail.com>
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Lancer
2026-03-10 16:58:21 +08:00
committed by GitHub
parent 2c2003158f
commit 8cd1de3354
5 changed files with 294 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files

57
python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py
View File

@@ -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()

25
python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py
View File

@@ -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(

118
python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py
View File

@@ -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()

18
python/sglang/multimodal_gen/test/cli/test_generate_common.py
View File

@@ -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",

132
python/sglang/multimodal_gen/test/cli/test_generate_i2i.py Normal file
View File

@@ -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()