From 9d45b9f0de16a1430c824093c17de9ccac47d25d Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Wed, 29 Oct 2025 02:23:49 +0200
Subject: [PATCH] Use comfy core Conv3D workaround for VAE rather than the fp32
 cast

---
 wanvideo/wan_video_vae.py | 29 +++--------------------------
 1 file changed, 3 insertions(+), 26 deletions(-)

diff --git a/wanvideo/wan_video_vae.py b/wanvideo/wan_video_vae.py
index cb2fd58..740822c 100644
--- a/wanvideo/wan_video_vae.py
+++ b/wanvideo/wan_video_vae.py
@@ -5,19 +5,11 @@ import torch.nn as nn
 import torch.nn.functional as F
 from tqdm import tqdm
 from comfy.utils import ProgressBar
+import comfy.ops
+ops = comfy.ops.disable_weight_init
 
 CACHE_T = 2
 
-# Workaround for increased memory usage in Conv3D with bfloat16 in torch 2.9.0 stable and up
-try:
-    torch_cudnn_bug = (
-        hasattr(torch.backends.cudnn, 'version') and
-        torch.backends.cudnn.version() >= 90800 and
-        torch.__version__ in ["2.9.0+cu126", "2.9.0+cu128", "2.9.0+cu130"] or torch.__version__.startswith("2.10.")
-    )
-except:
-    torch_cudnn_bug = False
-
 def check_is_instance(model, module_class):
     if isinstance(model, module_class):
         return True
@@ -26,7 +18,7 @@ def check_is_instance(model, module_class):
     return False
 
 
-class CausalConv3d(nn.Conv3d):
+class CausalConv3d(ops.Conv3d):
     """
     Causal 3d convolusion.
     """
@@ -45,21 +37,6 @@ class CausalConv3d(nn.Conv3d):
             padding[4] -= cache_x.shape[2]
         x = F.pad(x, padding)
 
-        # Convert to float32 only if this would trigger cuDNN bug
-        if (
-            torch_cudnn_bug and
-            x.dtype in (torch.bfloat16, torch.half) and
-            len(self.weight.shape) == 5 and
-            any(self.weight.shape[i] != 1 for i in range(2, 5))
-        ):
-            self.weight.data = self.weight.data.float()
-            if self.bias is not None:
-                self.bias.data = self.bias.data.float()
-            
-            result = super().forward(x.float())
-                
-            return result.to(x.dtype)
-
         return super().forward(x)