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]