anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Bugfix] Enable teacahce in QwenImageEditPlusPipeline (#379)

Browse Source 
Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com>
This commit is contained in:
Yuanheng Zhao authored and GitHub committed 2025-12-21 15:58:15 +08:00
1 parent 13584fd98e
commit 6b9944f11d
3 files changed
+33 -22
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/user_guide/acceleration/teacache.md
+1
View File
@@ -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
Vendored
+1
View File
@@ -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,
vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py
+31 -22
View File
@@ -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)