From f1bc6508ad63369ea2ac684a8ba929a9ab685c1e Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Fri, 17 Jul 2026 09:14:04 -0600 Subject: [PATCH] Fix issue with qwen image edit models. --- .../diffusion_models/qwen_image/qwen_image.py | 3 + .../qwen_image/qwen_image_edit.py | 2 - .../qwen_image/qwen_image_edit_plus.py | 2 - .../qwen_image/qwen_image_pipelines.py | 13 ---- .../util/qwen_vae_gradient_checkpointing.py | 71 +++++++++++++++++++ version.py | 2 +- 6 files changed, 75 insertions(+), 18 deletions(-) create mode 100644 toolkit/util/qwen_vae_gradient_checkpointing.py diff --git a/extensions_built_in/diffusion_models/qwen_image/qwen_image.py b/extensions_built_in/diffusion_models/qwen_image/qwen_image.py index 2e077fa..c03af90 100644 --- a/extensions_built_in/diffusion_models/qwen_image/qwen_image.py +++ b/extensions_built_in/diffusion_models/qwen_image/qwen_image.py @@ -30,6 +30,9 @@ from transformers import ( Qwen2VLProcessor, ) from tqdm import tqdm +from toolkit.util.qwen_vae_gradient_checkpointing import patch_qwen_vae_gradient_checkpointing + +patch_qwen_vae_gradient_checkpointing() if TYPE_CHECKING: from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO diff --git a/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit.py b/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit.py index 95880ac..efab9be 100644 --- a/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit.py +++ b/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit.py @@ -263,7 +263,6 @@ class QwenImageEditModel(QwenImageModel): prompt_embeds_mask = text_embeddings.attention_mask.to( self.device_torch, dtype=torch.int64 ) - txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() enc_hs = text_embeddings.text_embeds.to(self.device_torch, self.torch_dtype) noise_pred = self.transformer( @@ -273,7 +272,6 @@ class QwenImageEditModel(QwenImageModel): encoder_hidden_states=enc_hs, encoder_hidden_states_mask=prompt_embeds_mask.detach(), img_shapes=img_shapes, - txt_seq_lens=txt_seq_lens, return_dict=False, **kwargs, )[0] diff --git a/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit_plus.py b/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit_plus.py index 6273e1e..bb3a9c2 100644 --- a/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit_plus.py +++ b/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit_plus.py @@ -353,7 +353,6 @@ class QwenImageEditPlusModel(QwenImageModel): prompt_embeds_mask = text_embeddings.attention_mask.to( self.device_torch, dtype=torch.int64 ) - txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() enc_hs = text_embeddings.text_embeds.to(self.device_torch, self.torch_dtype) noise_pred = self.transformer( @@ -365,7 +364,6 @@ class QwenImageEditPlusModel(QwenImageModel): encoder_hidden_states=enc_hs.detach(), encoder_hidden_states_mask=prompt_embeds_mask.detach(), img_shapes=img_shapes, - txt_seq_lens=txt_seq_lens, return_dict=False, **kwargs, )[0] diff --git a/extensions_built_in/diffusion_models/qwen_image/qwen_image_pipelines.py b/extensions_built_in/diffusion_models/qwen_image/qwen_image_pipelines.py index dd6a9b1..addb962 100644 --- a/extensions_built_in/diffusion_models/qwen_image/qwen_image_pipelines.py +++ b/extensions_built_in/diffusion_models/qwen_image/qwen_image_pipelines.py @@ -230,17 +230,6 @@ class QwenImageEditPlusCustomPipeline(QwenImageEditPlusPipeline): if self.attention_kwargs is None: self._attention_kwargs = {} - txt_seq_lens = ( - prompt_embeds_mask.sum(dim=1).tolist() - if prompt_embeds_mask is not None - else None - ) - negative_txt_seq_lens = ( - negative_prompt_embeds_mask.sum(dim=1).tolist() - if negative_prompt_embeds_mask is not None - else None - ) - # 6. Denoising loop self.scheduler.set_begin_index(0) with self.progress_bar(total=num_inference_steps) as progress_bar: @@ -264,7 +253,6 @@ class QwenImageEditPlusCustomPipeline(QwenImageEditPlusPipeline): 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] @@ -279,7 +267,6 @@ class QwenImageEditPlusCustomPipeline(QwenImageEditPlusPipeline): 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] diff --git a/toolkit/util/qwen_vae_gradient_checkpointing.py b/toolkit/util/qwen_vae_gradient_checkpointing.py new file mode 100644 index 0000000..3e74c15 --- /dev/null +++ b/toolkit/util/qwen_vae_gradient_checkpointing.py @@ -0,0 +1,71 @@ +import torch +from diffusers.models.autoencoders.autoencoder_kl_qwenimage import ( + AutoencoderKLQwenImage, + QwenImageDecoder3d, + QwenImageEncoder3d, +) + +# diffusers removed gradient checkpointing from AutoencoderKLQwenImage +# (_supports_gradient_checkpointing = False and no checkpoint calls left in the +# forwards). This patches it back in so vae.enable_gradient_checkpointing() +# works again. + +_orig_encoder_forward = QwenImageEncoder3d.forward +_orig_decoder_forward = QwenImageDecoder3d.forward + +_patched = False + + +def _checkpoint(module, x): + return torch.utils.checkpoint.checkpoint(module, x, use_reentrant=False) + + +def _cache_is_fresh(feat_cache): + return feat_cache is None or all(c is None for c in feat_cache) + + +# The checkpointed paths run cache-free. feat_cache is mutated in place during +# the forward, so recomputing a block during backward would see end-state cache +# contents and produce garbage. With a fresh (all None) cache the cached and +# cache-free paths are numerically identical for the chunk being processed, so +# this is exact for single-frame (image) inputs. Later video chunks arrive with +# a populated cache and fall through to the original un-checkpointed forward. + + +def _encoder_forward(self, x, feat_cache=None, feat_idx=[0]): + if not (self.gradient_checkpointing and torch.is_grad_enabled() and _cache_is_fresh(feat_cache)): + return _orig_encoder_forward(self, x, feat_cache, feat_idx) + + x = self.conv_in(x) + for layer in self.down_blocks: + x = _checkpoint(layer, x) + x = _checkpoint(self.mid_block, x) + x = self.norm_out(x) + x = self.nonlinearity(x) + x = self.conv_out(x) + return x + + +def _decoder_forward(self, x, feat_cache=None, feat_idx=[0]): + if not (self.gradient_checkpointing and torch.is_grad_enabled() and _cache_is_fresh(feat_cache)): + return _orig_decoder_forward(self, x, feat_cache, feat_idx) + + x = self.conv_in(x) + x = _checkpoint(self.mid_block, x) + for up_block in self.up_blocks: + x = _checkpoint(up_block, x) + x = self.norm_out(x) + x = self.nonlinearity(x) + x = self.conv_out(x) + return x + + +def patch_qwen_vae_gradient_checkpointing(): + global _patched + if _patched: + return + _patched = True + + AutoencoderKLQwenImage._supports_gradient_checkpointing = True + QwenImageEncoder3d.forward = _encoder_forward + QwenImageDecoder3d.forward = _decoder_forward diff --git a/version.py b/version.py index 7ab44ab..9105263 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.10.27" +VERSION = "0.10.28"