[Bugfix] Qwen-Image use teachche serve will crash (#3450)

Signed-off-by: rongfu.leng <lenronfu@gmail.com>

tests/e2e/online_serving/test_qwen_image_edit_expansion.py

@@ -2,7 +2,7 @@
 Comprehensive tests of diffusion features that are available in online serving mode
 and are supported by the following models:
 - Qwen-Image-Edit: single image input
-- Qwen-Image-Edit-2509: two image inputs
+- Qwen-Image-Edit-2509: single image input and two image inputs
 """
 
 import pytest
@@ -14,6 +14,7 @@ from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHand
 pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
 
 EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting."
+SINGLE_EDIT_PROMPT_2509 = "Restyle this image into a Vincent van Gogh style impressionist painting."
 MULTI_EDIT_PROMPT = (
     "Transform the first image into a Dadaism collage art. "
     "Transform the second image into a Vincent van Gogh style painting. "
@@ -143,8 +144,43 @@ def test_qwen_image_edit(omni_server: OmniServer, openai_client: OpenAIClientHan
     _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2509"),
     indirect=True,
 )
-def test_qwen_image_edit_2509(omni_server: OmniServer, openai_client: OpenAIClientHandler):
-    """Test all diffusion features with Qwen-Image-Edit-2509 in regular end-user scenarios."""
+def test_qwen_image_edit_2509_single_image(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+    """Test Qwen-Image-Edit-2509 with a single image input.
+
+    Regression: with tea_cache enabled and zero_cond_t=True, the TeaCache
+    postprocess closure used the doubled temb (shape 2*batch) without halving
+    it, causing norm_out to broadcast and return noise_pred with shape
+    (2*batch, seq, ch). The scheduler step then silently expanded latents via
+    broadcasting, so at step 2 torch.cat([latents, image_latents], dim=1)
+    crashed with a batch size mismatch. Fixed in extractors.py.
+    """
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
+
+    messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=SINGLE_EDIT_PROMPT_2509)
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "extra_body": {
+            "height": 512,
+            "width": 512,
+            "num_inference_steps": 2,
+            "negative_prompt": NEGATIVE_PROMPT,
+            "true_cfg_scale": 4.0,
+            "seed": 42,
+        },
+    }
+
+    openai_client.send_diffusion_request(request_config)
+
+
+@pytest.mark.parametrize(
+    "omni_server",
+    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2509"),
+    indirect=True,
+)
+def test_qwen_image_edit_2509_two_images(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+    """Test Qwen-Image-Edit-2509 with two image inputs."""
     image_data_url_1 = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
     image_data_url_2 = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
 

vllm_omni/diffusion/cache/teacache/extractors.py

@@ -277,7 +277,8 @@ def extract_qwen_context(
 
     def postprocess(h):
         """Apply Qwen-specific output postprocessing."""
-        h = module.norm_out(h, temb)
+        t = temb.chunk(2, dim=0)[0] if module.zero_cond_t else temb
+        h = module.norm_out(h, t)
         output = module.proj_out(h)
         if not return_dict:
             return (output,)