Align Qwen-Image 2.1 preprocessing, CFG defaults and prefix cache (#7769)

Signed-off-by: Tianyu Guo <guoty9@mail2.sysu.edu.cn>
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>

examples/offline_inference/image_to_image/qwen_image_21_edit.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """
 Example script for image-conditioned generation with Qwen-Image 2.1.
@@ -168,7 +168,7 @@ def main():
     for image_path in args.image:
         if not os.path.exists(image_path):
             raise FileNotFoundError(f"Input image not found: {image_path}")
-        input_images.append(Image.open(image_path).convert("RGB"))
+        input_images.append(Image.open(image_path).convert("RGBA"))
 
     # Use single image or list based on number of inputs
     input_image = input_images[0] if len(input_images) == 1 else input_images
@@ -198,10 +198,7 @@ def main():
         print(f"  Number of input images: {len(input_image)}")
     else:
         print(f"  Input image size: {input_image.size}")
-    print(
-        f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, "
-        f"cfg_parallel_size={args.cfg_parallel_size}"
-    )
+    print(f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, cfg_parallel_size={args.cfg_parallel_size}")
     print(f"{'=' * 60}\n")
 
     generation_start = time.perf_counter()

vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py

@@ -6,7 +6,6 @@ import copy
 import json
 import logging
 import os
-import random
 from collections.abc import Iterable
 from typing import TYPE_CHECKING, Any, ClassVar
 
@@ -319,7 +318,7 @@ class QwenImage21Pipeline(
         # The prompt is built as a raw template string and passed straight to
         # `self.processor(text=..., images=...)`, rather than going through
         # `apply_chat_template`: the two tokenize differently and the checkpoint
-        # expects this one. The "Picture 1: ..." vision prefix only appears in
+        # expects this one. The "<image1>..." vision prefix only appears in
         # the image-conditioned template.
         self.prompt_template_t2i = (
             f"<|im_start|>system\n{self.sys_prompt}<|im_end|>\n"
@@ -328,10 +327,9 @@ class QwenImage21Pipeline(
         )
         self.prompt_template_ti2i = (
             f"<|im_start|>system\n{self.sys_prompt}<|im_end|>\n"
-            f"<|im_start|>user\nPicture 1: <|vision_start|><|image_pad|><|vision_end|>{{}}<|im_end|>\n"
+            f"<|im_start|>user\n<image1><|vision_start|><|image_pad|><|vision_end|>{{}}<|im_end|>\n"
             f"<|im_start|>assistant\n"
         )
-        self.ref_token_list = ["Picture ", "Image ", "图 ", "图片 "]
         # Number of leading system-role tokens to drop from the hidden states.
         # Derived from the tokenized system message rather than hardcoded, so it
         # tracks the processor's template.
@@ -470,6 +468,7 @@ class QwenImage21Pipeline(
         """
         dtype = self.text_encoder.dtype
         prompt = [prompt] if isinstance(prompt, str) else prompt
+        prompt = [" " if not p else p for p in prompt]
 
         has_images = images_per_prompt is not None and any(images_per_prompt)
         if not has_images:
@@ -488,15 +487,21 @@ class QwenImage21Pipeline(
                     txt.append(self.prompt_template_t2i.format(t))
                     continue
                 n_imgs = len(images)
-                replace = "Picture 1: <|vision_start|><|image_pad|><|vision_end|>"
+                replace = "<image1><|vision_start|><|image_pad|><|vision_end|>"
                 for i in range(2, n_imgs + 1):
-                    replace += f" Picture {i}: <|vision_start|><|image_pad|><|vision_end|>"
+                    replace += f" <image{i}><|vision_start|><|image_pad|><|vision_end|>"
                 template = self.prompt_template_ti2i.replace(
-                    "Picture 1: <|vision_start|><|image_pad|><|vision_end|>",
-                    replace.replace("Picture ", random.choice(self.ref_token_list)),
+                    "<image1><|vision_start|><|image_pad|><|vision_end|>", replace
                 )
                 txt.append(template.format(t))
-            condition_pil_list = [img for images in condition_pil_per_prompt for img in images]
+            condition_pil_list = []
+            for images in condition_pil_per_prompt:
+                for img in images:
+                    if img.mode == "RGBA":
+                        white = PIL.Image.new("RGB", img.size, (255, 255, 255))
+                        white.paste(img, mask=img.getchannel("A"))
+                        img = white
+                    condition_pil_list.append(img)
 
         # Validate only the user prompt contribution against the text budget;
         # image placeholder expansion happens later inside the processor.
@@ -521,7 +526,12 @@ class QwenImage21Pipeline(
             error_context="after applying the Qwen-Image 2.1 prompt template",
         )
 
-        processor_kwargs: dict[str, Any] = {"text": txt, "padding": True, "return_tensors": "pt"}
+        processor_kwargs: dict[str, Any] = {
+            "text": txt,
+            "padding": True,
+            "padding_side": "left",
+            "return_tensors": "pt",
+        }
         if has_images:
             processor_kwargs["images"] = condition_pil_list
         model_inputs = self.processor(**processor_kwargs).to(self.device)
@@ -976,7 +986,9 @@ class QwenImage21Pipeline(
         )
         generator = req.collate_request_generators(num_images_per_prompt, None)
         latents = req.collate_request_tensors("latents", None)
-        true_cfg_scale = common_sampling_params.true_cfg_scale or 4.0
+        true_cfg_scale = (
+            common_sampling_params.true_cfg_scale if common_sampling_params.true_cfg_scale is not None else 4.0
+        )
         output_type = common_sampling_params.output_type or "pil"
 
         # Batch homogeneity decides the attention path: a homogeneous batch
@@ -1062,7 +1074,7 @@ class QwenImage21Pipeline(
             sigmas=sampling.sigmas,
             num_images_per_prompt=num_images_per_prompt,
             generator=sampling.generator,
-            true_cfg_scale=sampling.true_cfg_scale or 4.0,
+            true_cfg_scale=sampling.true_cfg_scale if sampling.true_cfg_scale is not None else 4.0,
             max_sequence_length=sampling.max_sequence_length or self._max_length,
             per_request_images=per_request_images,
             attention_kwargs=kwargs.get("attention_kwargs"),

vllm_omni/diffusion/models/qwen_image_21/qwen_image_21_transformer.py

@@ -455,8 +455,9 @@ class QwenImage21Attention(nn.Module):
             branch_cache = kv_cache.setdefault(cache_branch, {})
             if cache_write_len is not None:
                 # Prefill: cache the timestep-independent prefix K/V for later denoising steps.
-                branch_cache["key"] = key[:, :cache_write_len].contiguous()
-                branch_cache["value"] = value[:, :cache_write_len].contiguous()
+                # Own the prefix storage instead of retaining the full prefill tensors.
+                branch_cache["key"] = key[:, :cache_write_len].clone()
+                branch_cache["value"] = value[:, :cache_write_len].clone()
             else:
                 cached_key = branch_cache["key"]
                 cached_value = branch_cache["value"]