[Feature] Enable TeaCache in QwenImageEditPipeline (#304)

Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com>

docs/user_guide/acceleration/teacache.md

@@ -87,13 +87,28 @@ cache_config={"rel_l1_thresh": 0.1}  # More conservative caching
 
 ## Supported Models
 
-Currently supported models:
+<style>
+th {
+  white-space: nowrap;
+  min-width: 0 !important;
+}
+</style>
 
-- **QwenImagePipeline**
-  - Model identifier: `"Qwen/Qwen-Image"`
-  - Example: `model="Qwen/Qwen-Image"`
+| Architecture | Models | Example HF Models |
+|--------------|--------|-------------------|
+| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
+| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` |
 
-Future support (extractors in development):
+### Coming Soon
 
-- **FluxPipeline** - Coming soon
-- **CogVideoXPipeline** - Coming soon
+<style>
+th {
+  white-space: nowrap;
+  min-width: 0 !important;
+}
+</style>
+
+| Architecture | Models | Example HF Models |
+|--------------|--------|-------------------|
+| `FluxPipeline` | Flux | - |
+| `CogVideoXPipeline` | CogVideoX | - |

examples/offline_inference/image_to_image/image_edit.py

@@ -144,7 +144,14 @@ def main():
             "scm_steps_policy": "dynamic",  # SCM steps policy: "dynamic" or "static"
         }
     elif args.cache_backend == "tea_cache":
-        raise ValueError("TeaCache is not supported for image-to-image generation.")
+        # TeaCache configuration
+        # All parameters marked with [tea_cache only] in DiffusionCacheConfig
+        cache_config = {
+            # TeaCache parameters [tea_cache only]
+            "rel_l1_thresh": 0.2,  # Threshold for accumulated relative L1 distance
+            # Note: coefficients will use model-specific defaults based on model_type
+            #       (e.g., QwenImagePipeline or FluxPipeline)
+        }
 
     # Initialize Omni with QwenImageEditPipeline
     omni = Omni(

vllm_omni/diffusion/cache/teacache/extractors.py

@@ -261,6 +261,7 @@ def extract_qwen_context(
 # to ensure consistent lookup without substring matching.
 EXTRACTOR_REGISTRY: dict[str, Callable] = {
     "QwenImagePipeline": extract_qwen_context,
+    "QwenImageEditPipeline": extract_qwen_context,
     # Future models:
     # "FluxPipeline": extract_flux_context,
     # "CogVideoXPipeline": extract_cogvideox_context,

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py

@@ -230,6 +230,9 @@ class QwenImageEditPipeline(
             model, subfolder="processor", local_files_only=local_files_only
         )
 
+        # Initialize cache backend to None (will be set by worker if needed)
+        self._cache_backend = None
+
         self.stage = None
 
         self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8
@@ -601,31 +604,43 @@ class QwenImageEditPipeline(
             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]
+            # Forward pass for positive prompt (or unconditional if no CFG)
+            # cache_branch is passed to hook for CFG-aware state management
+            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)]
 
+            # Forward pass for negative prompt (CFG)
+            # cache_branch is passed to hook for CFG-aware state management
             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)