diff --git a/docs/user_guide/acceleration/teacache.md b/docs/user_guide/acceleration/teacache.md
index 5df0a4e1a..ef7af50c3 100644
--- a/docs/user_guide/acceleration/teacache.md
+++ b/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:
+
-- **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
+
+
+| Architecture | Models | Example HF Models |
+|--------------|--------|-------------------|
+| `FluxPipeline` | Flux | - |
+| `CogVideoXPipeline` | CogVideoX | - |
diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py
index 33a95454d..9547a3cd0 100644
--- a/examples/offline_inference/image_to_image/image_edit.py
+++ b/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(
diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py
index 49bcb76ca..0bb10ff82 100644
--- a/vllm_omni/diffusion/cache/teacache/extractors.py
+++ b/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,
diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py
index c60623d81..8605e01a5 100644
--- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py
+++ b/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)