From 886c2aec57d4ef10a6b67d2f4e8f95e6067866fb Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Fri, 10 Jul 2026 09:42:52 -0600
Subject: [PATCH] Allow for vae tiling onle without low vram on wan models with
 a model kwarg

---
 .../diffusion_models/wan22/wan22_14b_model.py    |  9 +++++++++
 .../diffusion_models/wan22/wan22_5b_model.py     |  4 ++--
 toolkit/models/wan21/wan21.py                    | 16 ++++++++++++++++
 3 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py b/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py
index be72fb5..37bcb7a 100644
--- a/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py
+++ b/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py
@@ -555,6 +555,11 @@ class Wan2214bModel(Wan21):
     ):
         # reactivate progress bar since this is slooooow
         pipeline.set_progress_bar_config(disable=False)
+
+        if self.use_vae_tiling:
+            # set vae to tile decode
+            pipeline.vae.enable_tiling()
+
         # todo, figure out how to do video
         output = pipeline(
             prompt_embeds=conditional_embeds.text_embeds.to(
@@ -573,6 +578,10 @@ class Wan2214bModel(Wan21):
             **extra
         )[0]
 
+        if self.use_vae_tiling:
+            # restore no tiling
+            pipeline.vae.disable_tiling()
+
         # shape = [1, frames, channels, height, width]
         batch_item = output[0]  # list of pil images
         if gen_config.num_frames > 1:
diff --git a/extensions_built_in/diffusion_models/wan22/wan22_5b_model.py b/extensions_built_in/diffusion_models/wan22/wan22_5b_model.py
index f231d74..5d8e48f 100644
--- a/extensions_built_in/diffusion_models/wan22/wan22_5b_model.py
+++ b/extensions_built_in/diffusion_models/wan22/wan22_5b_model.py
@@ -210,7 +210,7 @@ class Wan225bModel(Wan21):
                 latent_model_input=latents, first_frame=first_frame_n1p1, vae=self.vae
             )
 
-        if self.model_config.low_vram:
+        if self.use_vae_tiling:
             # set vae to tile decode
             pipeline.vae.enable_tiling()
 
@@ -234,7 +234,7 @@ class Wan225bModel(Wan21):
             **extra,
         )[0]
 
-        if self.model_config.low_vram:
+        if self.use_vae_tiling:
             # restore no tiling
             pipeline.vae.disable_tiling()
 
diff --git a/toolkit/models/wan21/wan21.py b/toolkit/models/wan21/wan21.py
index 8519b7d..3f7f953 100644
--- a/toolkit/models/wan21/wan21.py
+++ b/toolkit/models/wan21/wan21.py
@@ -520,6 +520,13 @@ class Wan21(BaseModel):
 
         return pipeline
 
+    @property
+    def use_vae_tiling(self):
+        # tile the vae decode when sampling if in low vram or explicitly enabled
+        return self.model_config.low_vram or self.model_config.model_kwargs.get(
+            "vae_tiling", False
+        )
+
     def generate_single_image(
         self,
         pipeline: WanPipeline,
@@ -532,6 +539,11 @@ class Wan21(BaseModel):
         # reactivate progress bar since this is slooooow
         pipeline.set_progress_bar_config(disable=False)
         pipeline = pipeline.to(self.device_torch)
+
+        if self.use_vae_tiling:
+            # set vae to tile decode
+            pipeline.vae.enable_tiling()
+
         # todo, figure out how to do video
         output = pipeline(
             prompt_embeds=conditional_embeds.text_embeds.to(
@@ -550,6 +562,10 @@ class Wan21(BaseModel):
             **extra
         )[0]
 
+        if self.use_vae_tiling:
+            # restore no tiling
+            pipeline.vae.disable_tiling()
+
         # shape = [1, frames, channels, height, width]
         batch_item = output[0]  # list of pil images
         if gen_config.num_frames > 1: