[diffusion] model: support mutli-image input and qwen-image-edit-2509 (#15005)

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

@@ -208,7 +208,10 @@ class PipelineConfig:
             (target_width, target_height), PIL.Image.Resampling.LANCZOS
         ), (target_width, target_height)
 
-    def prepare_image_processor_kwargs(self, batch):
+    def prepare_calculated_size(self, image):
+        return self.calculate_condition_image_size(image, image.width, image.height)
+
+    def prepare_image_processor_kwargs(self, batch, neg=False):
         return {}
 
     def postprocess_image_latent(self, latent_condition, batch):
@@ -297,6 +300,9 @@ class PipelineConfig:
         latents = sequence_model_parallel_all_gather(latents, dim=2)
         return latents
 
+    def preprocess_vae_image(self, batch, vae_image_processor):
+        pass
+
     def shard_latents_for_sp(self, batch, latents):
         # general logic for video models
         sp_world_size, rank_in_sp_group = get_sp_world_size(), get_sp_parallel_rank()

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

@@ -107,8 +107,9 @@ class QwenImagePipelineConfig(ImagePipelineConfig):
     def prepare_sigmas(self, sigmas, num_inference_steps):
         return self._prepare_sigmas(sigmas, num_inference_steps)
 
-    def prepare_image_processor_kwargs(self, batch):
-        if batch.prompt:
+    def prepare_image_processor_kwargs(self, batch, neg=False):
+        prompt = batch.prompt if not neg else batch.negative_prompt
+        if prompt:
             prompt_template_encode = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n<|im_start|>assistant\n"
             txt = prompt_template_encode.format(batch.prompt)
             return dict(text=[txt], padding=True)
@@ -328,3 +329,137 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig):
         # remove noise over input image
         noise = noise[:, : latents.size(1)]
         return noise
+
+
+CONDITION_IMAGE_SIZE = 384 * 384
+VAE_IMAGE_SIZE = 1024 * 1024
+
+
+@dataclass
+class QwenImageEditPlusPipelineConfig(QwenImageEditPipelineConfig):
+    task_type: ModelTaskType = ModelTaskType.I2I
+
+    def _get_condition_image_sizes(self, batch) -> list[tuple[int, int]]:
+        image = batch.condition_image
+        if not isinstance(image, list):
+            image = [image]
+
+        condition_image_sizes = []
+        for img in image:
+            image_width, image_height = img.size
+            edit_width, edit_height, _ = calculate_dimensions(
+                VAE_IMAGE_SIZE, image_width / image_height
+            )
+            condition_image_sizes.append((edit_width, edit_height))
+
+        return condition_image_sizes
+
+    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
+
+        prompt_template_encode = (
+            "<|im_start|>system\nDescribe the key features of the input image "
+            "(color, shape, size, texture, objects, background), then explain how "
+            "the user's text instruction should alter or modify the image. Generate "
+            "a new image that meets the user's requirements while maintaining "
+            "consistency with the original input where appropriate.<|im_end|>\n"
+            "<|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)
+
+    def prepare_calculated_size(self, image):
+        return self.calculate_vae_image_size(image, image.width, image.height)
+
+    def resize_condition_image(self, images, target_width, target_height):
+        if not isinstance(images, list):
+            images = [images]
+        new_images = []
+        for img, width, height in zip(images, target_width, target_height):
+            new_images.append(resize(img, height, width, resize_mode="default"))
+        return new_images
+
+    def calculate_condition_image_size(self, image, width, height) -> tuple[int, int]:
+        calculated_width, calculated_height, _ = calculate_dimensions(
+            CONDITION_IMAGE_SIZE, width / height
+        )
+        return calculated_width, calculated_height
+
+    def calculate_vae_image_size(self, image, width, height) -> tuple[int, int]:
+        calculated_width, calculated_height, _ = calculate_dimensions(
+            VAE_IMAGE_SIZE, width / height
+        )
+        return calculated_width, calculated_height
+
+    def preprocess_vae_image(self, batch, vae_image_processor):
+        if not isinstance(batch.condition_image, list):
+            batch.condition_image = [batch.condition_image]
+        new_images = []
+        vae_image_sizes = []
+        for img in batch.condition_image:
+            width, height = self.calculate_vae_image_size(img, img.width, img.height)
+            new_images.append(vae_image_processor.preprocess(img, height, width))
+            vae_image_sizes.append((width, height))
+        batch.vae_image = new_images
+        batch.vae_image_sizes = vae_image_sizes
+        return batch
+
+    def _prepare_edit_cond_kwargs(
+        self, batch, prompt_embeds, rotary_emb, device, dtype
+    ):
+        batch_size = batch.latents.shape[0]
+        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()
+
+        img_shapes = [
+            [
+                (1, height // vae_scale_factor // 2, width // vae_scale_factor // 2),
+                *[
+                    (
+                        1,
+                        vae_height // vae_scale_factor // 2,
+                        vae_width // vae_scale_factor // 2,
+                    )
+                    for vae_width, vae_height in batch.vae_image_sizes
+                ],
+            ],
+        ] * batch_size
+        txt_seq_lens = [prompt_embeds[0].shape[1]]
+
+        (img_cos, img_sin), (txt_cos, txt_sin) = (
+            QwenImageEditPlusPipelineConfig.get_freqs_cis(
+                img_shapes, txt_seq_lens, rotary_emb, device, dtype
+            )
+        )
+
+        # perform sp shard on noisy image tokens
+        noisy_img_seq_len = (
+            1 * (height // vae_scale_factor // 2) * (width // vae_scale_factor // 2)
+        )
+        noisy_img_cos = shard_rotary_emb_for_sp(img_cos[:noisy_img_seq_len, :])
+        noisy_img_sin = shard_rotary_emb_for_sp(img_sin[:noisy_img_seq_len, :])
+
+        # concat back the img_cos for input image (since it is not sp-shared later)
+        img_cos = torch.cat([noisy_img_cos, img_cos[noisy_img_seq_len:, :]], dim=0).to(
+            device=device
+        )
+        img_sin = torch.cat([noisy_img_sin, img_sin[noisy_img_seq_len:, :]], dim=0).to(
+            device=device
+        )
+
+        return {
+            "txt_seq_lens": txt_seq_lens,
+            "freqs_cis": ((img_cos, img_sin), (txt_cos, txt_sin)),
+        }

python/sglang/multimodal_gen/configs/sample/qwenimage.py

@@ -13,3 +13,11 @@ class QwenImageSamplingParams(SamplingParams):
     # Denoising stage
     guidance_scale: float = 4.0
     num_inference_steps: int = 50
+
+
+@dataclass
+class QwenImageEditPlusSamplingParams(QwenImageSamplingParams):
+    # Denoising stage
+    guidance_scale: float = 4.0
+    # true_cfg_scale: float = 4.0
+    num_inference_steps: int = 40