From 963a9f42b2444955ad8243cd4f96e245961ae0d2 Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Tue, 5 May 2026 11:30:16 -0600
Subject: [PATCH] Add decode latent to wan 2.1 models. Add gradinet
 checkpointing to wan vae.

---
 toolkit/models/wan21/autoencoder_kl_wan.py | 31 ++++++++++++++--------
 toolkit/models/wan21/wan21.py              | 25 +++++++++++++++++
 2 files changed, 45 insertions(+), 11 deletions(-)

diff --git a/toolkit/models/wan21/autoencoder_kl_wan.py b/toolkit/models/wan21/autoencoder_kl_wan.py
index 3a96610..8761eea 100644
--- a/toolkit/models/wan21/autoencoder_kl_wan.py
+++ b/toolkit/models/wan21/autoencoder_kl_wan.py
@@ -450,8 +450,6 @@ class WanMidBlock(nn.Module):
         self.attentions = nn.ModuleList(attentions)
         self.resnets = nn.ModuleList(resnets)
 
-        self.gradient_checkpointing = False
-
     def forward(self, x, feat_cache=None, feat_idx=[0]):
         # First residual block
         x = self.resnets[0](x, feat_cache, feat_idx)
@@ -602,15 +600,22 @@ class WanEncoder3d(nn.Module):
         else:
             x = self.conv_in(x)
 
+        use_ckpt = torch.is_grad_enabled() and self.gradient_checkpointing and feat_cache is None
+
         ## downsamples
         for layer in self.down_blocks:
-            if feat_cache is not None:
+            if use_ckpt:
+                x = self._gradient_checkpointing_func(layer, x)
+            elif feat_cache is not None:
                 x = layer(x, feat_cache, feat_idx)
             else:
                 x = layer(x)
 
         ## middle
-        x = self.mid_block(x, feat_cache, feat_idx)
+        if use_ckpt:
+            x = self._gradient_checkpointing_func(self.mid_block, x)
+        else:
+            x = self.mid_block(x, feat_cache, feat_idx)
 
         ## head
         x = self.norm_out(x)
@@ -682,8 +687,6 @@ class WanResidualUpBlock(nn.Module):
         else:
             self.upsampler = None
 
-        self.gradient_checkpointing = False
-
     def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False):
         """
         Forward pass through the upsampling block.
@@ -756,8 +759,6 @@ class WanUpBlock(nn.Module):
         if upsample_mode is not None:
             self.upsamplers = nn.ModuleList([WanResample(out_dim, mode=upsample_mode)])
 
-        self.gradient_checkpointing = False
-
     def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=None):
         """
         Forward pass through the upsampling block.
@@ -889,12 +890,20 @@ class WanDecoder3d(nn.Module):
         else:
             x = self.conv_in(x)
 
+        use_ckpt = torch.is_grad_enabled() and self.gradient_checkpointing and feat_cache is None
+
         ## middle
-        x = self.mid_block(x, feat_cache, feat_idx)
+        if use_ckpt:
+            x = self._gradient_checkpointing_func(self.mid_block, x)
+        else:
+            x = self.mid_block(x, feat_cache, feat_idx)
 
         ## upsamples
         for up_block in self.up_blocks:
-            x = up_block(x, feat_cache, feat_idx, first_chunk = first_chunk)
+            if use_ckpt:
+                x = self._gradient_checkpointing_func(up_block, x, None, [0], first_chunk)
+            else:
+                x = up_block(x, feat_cache, feat_idx, first_chunk = first_chunk)
 
         ## head
         x = self.norm_out(x)
@@ -961,7 +970,7 @@ class AutoencoderKLWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
     for all models (such as downloading or saving).
     """
 
-    _supports_gradient_checkpointing = False
+    _supports_gradient_checkpointing = True
 
     @register_to_config
     def __init__(
diff --git a/toolkit/models/wan21/wan21.py b/toolkit/models/wan21/wan21.py
index 7e14bdc..8519b7d 100644
--- a/toolkit/models/wan21/wan21.py
+++ b/toolkit/models/wan21/wan21.py
@@ -654,6 +654,31 @@ class Wan21(BaseModel):
         latents = (latents - latents_mean) * latents_std
 
         return latents.to(device, dtype=dtype)
+    
+    def decode_latents(self, latents: torch.Tensor, device=None, dtype=None):
+        if device is None:
+            device = self.vae_device_torch
+        if dtype is None:
+            dtype = self.vae_torch_dtype
+
+        if self.vae.device == torch.device('cpu'):
+            self.vae.to(device)
+
+        latents = latents.to(device, dtype=dtype)
+
+        latents_mean = (
+            torch.tensor(self.vae.config.latents_mean)
+            .view(1, self.vae.config.z_dim, 1, 1, 1)
+            .to(latents.device, latents.dtype)
+        )
+        latents_std = torch.tensor(self.vae.config.latents_std).view(
+            1, self.vae.config.z_dim, 1, 1, 1
+        ).to(latents.device, latents.dtype)
+        latents = latents * latents_std + latents_mean
+
+        images = self.vae.decode(latents).sample
+
+        return images.to(device, dtype=dtype)
 
     def get_model_has_grad(self):
         return False