[Bugfix] Enable teacahce in QwenImageEditPlusPipeline (vllm-project#379)

Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com>
Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com>

Author: yuanheng-zhao

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
 

vllm_omni/diffusion/cache/teacache/extractors.py

@@ -262,6 +262,7 @@ def postprocess(h):
 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,

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py

@@ -194,6 +194,7 @@ def __init__(
         )
 
         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 @@ def diffuse(
             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)