[Bugfix] Enable teacahce in QwenImageEditPlusPipeline

[yuanheng-zhao]

PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.

Purpose

Enable TeaCache in QwenImageEditPlusPipeline

Test Plan

Offline inference with running the example script examples/offline_inference/image_to_image/image_edit.py

python image_edit.py \
  --model Qwen/Qwen-Image-Edit-2509 \
  --image /path/to/coffee.png /path/to/cupcake.png \
  --prompt "Put these two items on the same table" \
  --output cupcake_edit.png \
  --num_inference_steps 50 \
  --cfg_scale 4.0 \
  --guidance_scale 1.0 \
  --cache_backend tea_cache

python image_edit.py \
  --model Qwen/Qwen-Image-Edit-2509 \
  --image /path/to/coffee.png /path/to/cupcake.png \
  --prompt "Put these two items on the same table" \
  --output cupcake_edit_cachedit.png \
  --num_inference_steps 50 \
  --cfg_scale 4.0 \
  --guidance_scale 1.0 \
  --cache_backend cache_dit

Test Result

Tested on H100.

Total Generation Time

Edited Image (no cache)	Edited Image w/ TeaCache	Edited Image w/ CacheDiT
81.6388 sec	61.2393 sec	38.0708 sec

Input Images

1st Input Image	2nd Input Image

Edited Images

Edited Image	Edited Image w/ TeaCache	Edited Image w/ CacheDiT

Note: During testing the following error happened because of the warmup dummy run does not include an image path into the request. This won't affect the execution of actual request though.

INFO 12-19 08:28:23 [diffusion_engine.py:173] dummy run to warm up the model
ERROR 12-19 08:28:23 [gpu_worker.py:209] Error executing forward in event loop: Image is required for QwenImageEditPlusPipeline
...
ERROR 12-19 08:28:23 [gpu_worker.py:209]     raise ValueError("Image is required for QwenImageEditPlusPipeline")
ERROR 12-19 08:28:23 [gpu_worker.py:209] ValueError: Image is required for QwenImageEditPlusPipeline

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft.

BEFORE SUBMITTING, PLEASE READ https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md (anything written below this line will be removed by GitHub Actions)

[ZJY0516]

ERROR 12-19 08:28:23 [gpu_worker.py:209] Error executing forward in event loop: Image is required for QwenImageEditPlusPipeline

good catch. I'll fix it later

[SamitHuang]

lgtm

[hsliuustc0106]

lgtm, one question: does it apply to other qwen-image models as well or we need another a PR?

[yuanheng-zhao]

lgtm, one question: does it apply to other qwen-image models as well or we need another a PR?

Hey @hsliuustc0106 , TeaCache has been enabled for the following Qwen-Image family pipelines, with the same extractor extract_qwen_context registered, as extract_qwen_context is specifically targeting QwenImageTransformer2DModel:

• QwenImagePipeline
• QwenImageEditPipeline
• QwenImageEditPlusPipeline

It's expected to work for other Qwen family models which applies the same transformer forward flow. However, there might be some arg diff to adapt and test, like extra args imported into model and transformer block forward methods by QwenImageLayeredPipeline recently.

Another note: As there exist many duplicate / similar code patterns for Qwen-Image-family pipelines, we might want to re-use some code and reduce redundancy, by applying a CacheMixin or a QwenImageBasePipeline (or some cleaner pattern), in the future.

[ZJY0516]

Another note: As there exist many duplicate / similar code patterns for Qwen-Image-family pipelines, we might want to re-use some code and reduce redundancy, by applying a CacheMixin or a QwenImageBasePipeline (or some cleaner pattern), in the future.

PTAL #395

[ZJY0516]

It's expected to work for other Qwen family models which applies the same transformer forward flow. However, there might be some arg diff to adapt and test, like extra args imported into model and transformer block forward methods by QwenImageLayeredPipeline recently.

Feel free to open a pr

[yuanheng-zhao]

It's expected to work for other Qwen family models which applies the same transformer forward flow. However, there might be some arg diff to adapt and test, like extra args imported into model and transformer block forward methods by QwenImageLayeredPipeline recently.

Feel free to open a pr

Did a bit cleanup in #416