Fix issue with qwen image edit models.

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

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]

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]

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]

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

version.py

@@ -1 +1 @@
-VERSION = "0.10.27"
+VERSION = "0.10.28"