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)