From 0323a9d2c74231731ca0e75304057cd574f5f675 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Thu, 29 May 2025 18:37:13 +0300
Subject: [PATCH] Fix TAEW decoding

---
 nodes.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/nodes.py b/nodes.py
index b317921..b21d28f 100644
--- a/nodes.py
+++ b/nodes.py
@@ -3487,14 +3487,16 @@ class WanVideoDecode:
 
         #if is_looped:
         #   latents = torch.cat([latents[:, :, :warmup_latent_count],latents], dim=2)
-
-        if isinstance(vae, TAEHV):            
+        if type(vae).__name__ == "TAEHV":      
             images = vae.decode_video(latents.permute(0, 2, 1, 3, 4))[0].permute(1, 0, 2, 3)
+            images = torch.clamp(images, 0.0, 1.0)
+            images = images.permute(1, 2, 3, 0).cpu().float()
+            return (images,)
         else:
             if end_image is not None:
                 enable_vae_tiling = False
             images = vae.decode(latents, device=device, end_=(end_image is not None), tiled=enable_vae_tiling, tile_size=(tile_x//8, tile_y//8), tile_stride=(tile_stride_x//8, tile_stride_y//8))[0]
-        vae.model.clear_cache()
+            vae.model.clear_cache()
 
         images = (images - images.min()) / (images.max() - images.min())