From fe82487187bb4c0e61e885d421fc6ba9563073de Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Fri, 10 Jul 2026 09:07:29 -0600
Subject: [PATCH] Add tiling on vae decode for qwen image models when low vram
 flag is on

---
 .../diffusion_models/qwen_image/qwen_image.py          | 10 ++++++++++
 .../diffusion_models/qwen_image/qwen_image_edit.py     |  9 +++++++++
 .../qwen_image/qwen_image_edit_plus.py                 |  9 +++++++++
 3 files changed, 28 insertions(+)

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 648f22b..2e077fa 100644
--- a/extensions_built_in/diffusion_models/qwen_image/qwen_image.py
+++ b/extensions_built_in/diffusion_models/qwen_image/qwen_image.py
@@ -284,6 +284,11 @@ class QwenImageModel(BaseModel):
         sc = self.get_bucket_divisibility()
         gen_config.width = int(gen_config.width // sc * sc)
         gen_config.height = int(gen_config.height // sc * sc)
+
+        if self.model_config.low_vram:
+            # set vae to tile decode
+            pipeline.vae.enable_tiling()
+
         img = pipeline(
             prompt_embeds=conditional_embeds.text_embeds,
             prompt_embeds_mask=conditional_embeds.attention_mask.to(
@@ -302,6 +307,11 @@ class QwenImageModel(BaseModel):
             callback_on_step_end=callback_on_step_end,
             **extra,
         ).images[0]
+
+        if self.model_config.low_vram:
+            # restore no tiling
+            pipeline.vae.disable_tiling()
+
         return img
 
     def get_noise_prediction(
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 6061889..95880ac 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
@@ -115,6 +115,10 @@ class QwenImageEditModel(QwenImageModel):
 
             return {"latents": latents}
 
+        if self.model_config.low_vram:
+            # set vae to tile decode
+            pipeline.vae.enable_tiling()
+
         img = pipeline(
             image=control_img,
             prompt_embeds=conditional_embeds.text_embeds,
@@ -134,6 +138,11 @@ class QwenImageEditModel(QwenImageModel):
             callback_on_step_end=callback_on_step_end,
             **extra,
         ).images[0]
+
+        if self.model_config.low_vram:
+            # restore no tiling
+            pipeline.vae.disable_tiling()
+
         return img
 
     def condition_noisy_latents(
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 1f5a376..6273e1e 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
@@ -133,6 +133,10 @@ class QwenImageEditPlusModel(QwenImageModel):
 
             return {"latents": latents}
 
+        if self.model_config.low_vram:
+            # set vae to tile decode
+            pipeline.vae.enable_tiling()
+
         img = pipeline(
             image=control_img_list,
             prompt_embeds=conditional_embeds.text_embeds,
@@ -153,6 +157,11 @@ class QwenImageEditPlusModel(QwenImageModel):
             do_cfg_norm=gen_config.do_cfg_norm,
             **extra,
         ).images[0]
+
+        if self.model_config.low_vram:
+            # restore no tiling
+            pipeline.vae.disable_tiling()
+
         return img
 
     def condition_noisy_latents(