From b4ab37da22e77a112e6f6e085937a4ea66ed6da9 Mon Sep 17 00:00:00 2001 From: "rongfu.leng" Date: Sat, 9 May 2026 06:41:59 +0800 Subject: [PATCH] [Bugfix] Qwen-Image use teachche serve will crash (#3450) Signed-off-by: rongfu.leng --- .../test_qwen_image_edit_expansion.py | 42 +++++++++++++++++-- .../diffusion/cache/teacache/extractors.py | 3 +- 2 files changed, 41 insertions(+), 4 deletions(-) diff --git a/tests/e2e/online_serving/test_qwen_image_edit_expansion.py b/tests/e2e/online_serving/test_qwen_image_edit_expansion.py index bb8bf7dca..9b30ef4eb 100644 --- a/tests/e2e/online_serving/test_qwen_image_edit_expansion.py +++ b/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']}" diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py index 6c5acdbf3..5b7ab14e5 100644 --- a/vllm_omni/diffusion/cache/teacache/extractors.py +++ b/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,)