From 6b9944f11d44d92b322f44bb0d5046c40cf0f12e Mon Sep 17 00:00:00 2001 From: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com> Date: Sun, 21 Dec 2025 15:58:15 +0800 Subject: [PATCH] [Bugfix] Enable teacahce in QwenImageEditPlusPipeline (#379) Signed-off-by: yuanheng --- docs/user_guide/acceleration/teacache.md | 1 + .../diffusion/cache/teacache/extractors.py | 1 + .../pipeline_qwen_image_edit_plus.py | 53 +++++++++++-------- 3 files changed, 33 insertions(+), 22 deletions(-) diff --git a/docs/user_guide/acceleration/teacache.md b/docs/user_guide/acceleration/teacache.md index ab8958929..641a43048 100644 --- a/docs/user_guide/acceleration/teacache.md +++ b/docs/user_guide/acceleration/teacache.md @@ -108,6 +108,7 @@ th { |--------------|--------|-------------------| | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | | `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` | +| `QwenImageEditPlusPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit-2509` | ### Coming Soon diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py index 0bb10ff82..8047e964e 100644 --- a/vllm_omni/diffusion/cache/teacache/extractors.py +++ b/vllm_omni/diffusion/cache/teacache/extractors.py @@ -262,6 +262,7 @@ def extract_qwen_context( EXTRACTOR_REGISTRY: dict[str, Callable] = { "QwenImagePipeline": extract_qwen_context, "QwenImageEditPipeline": extract_qwen_context, + "QwenImageEditPlusPipeline": extract_qwen_context, # Future models: # "FluxPipeline": extract_flux_context, # "CogVideoXPipeline": extract_cogvideox_context, diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py index d3b9bfe94..0042484a7 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py @@ -194,6 +194,7 @@ class QwenImageEditPlusPipeline( ) self.stage = None + self._cache_backend = None self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8 self.latent_channels = self.vae.config.z_dim if getattr(self, "vae", None) else 16 @@ -543,31 +544,39 @@ class QwenImageEditPlusPipeline( if image_latents is not None: latent_model_input = torch.cat([latents, image_latents], dim=1) - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep / 1000, - guidance=guidance, - encoder_hidden_states_mask=prompt_embeds_mask, - encoder_hidden_states=prompt_embeds, - img_shapes=img_shapes, - txt_seq_lens=txt_seq_lens, - attention_kwargs=self.attention_kwargs, - return_dict=False, - )[0] + transformer_kwargs = { + "hidden_states": latent_model_input, + "timestep": timestep / 1000, + "guidance": guidance, + "encoder_hidden_states_mask": prompt_embeds_mask, + "encoder_hidden_states": prompt_embeds, + "img_shapes": img_shapes, + "txt_seq_lens": txt_seq_lens, + "attention_kwargs": self.attention_kwargs, + "return_dict": False, + } + if self._cache_backend is not None: + transformer_kwargs["cache_branch"] = "positive" + + noise_pred = self.transformer(**transformer_kwargs)[0] noise_pred = noise_pred[:, : latents.size(1)] if do_true_cfg: - neg_noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep / 1000, - guidance=guidance, - encoder_hidden_states_mask=negative_prompt_embeds_mask, - encoder_hidden_states=negative_prompt_embeds, - img_shapes=img_shapes, - txt_seq_lens=negative_txt_seq_lens, - attention_kwargs=self.attention_kwargs, - return_dict=False, - )[0] + neg_transformer_kwargs = { + "hidden_states": latent_model_input, + "timestep": timestep / 1000, + "guidance": guidance, + "encoder_hidden_states_mask": negative_prompt_embeds_mask, + "encoder_hidden_states": negative_prompt_embeds, + "img_shapes": img_shapes, + "txt_seq_lens": negative_txt_seq_lens, + "attention_kwargs": self.attention_kwargs, + "return_dict": False, + } + if self._cache_backend is not None: + neg_transformer_kwargs["cache_branch"] = "negative" + + neg_noise_pred = self.transformer(**neg_transformer_kwargs)[0] neg_noise_pred = neg_noise_pred[:, : latents.size(1)] comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred) cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)