From 9e957121bea76b07b56fe6df33b5ebba2a553295 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Thu, 27 Feb 2025 11:13:30 +0200
Subject: [PATCH] better preview

---
 nodes.py                      | 13 +++++++++----
 wanvideo/modules/attention.py |  8 ++++----
 wanvideo/modules/model.py     |  3 ++-
 3 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/nodes.py b/nodes.py
index c6de497..0468562 100644
--- a/nodes.py
+++ b/nodes.py
@@ -102,14 +102,19 @@ class WanVideoModel(comfy.model_base.BaseModel):
     def __setitem__(self, k, v):
         self.pipeline[k] = v
 
-from comfy.latent_formats import LatentFormat
-
+try:
+    from comfy.latent_formats import Wan21
+    latent_format = Wan21
+except: #for backwards compatibility
+    log.warning("Wan21 latent format not found, update ComfyUI for better livepreview")
+    from comfy.latent_formats import HunyuanVideo
+    latent_format = HunyuanVideo
 
 class WanVideoModelConfig:
     def __init__(self, dtype):
         self.unet_config = {}
         self.unet_extra_config = {}
-        self.latent_format = comfy.latent_formats.HunyuanVideo #todo better values
+        self.latent_format = latent_format
         self.latent_format.latent_channels = 16
         self.manual_cast_dtype = dtype
         self.sampling_settings = {"multiplier": 1.0}
@@ -960,7 +965,7 @@ class WanVideoSampler:
 
         mm.soft_empty_cache()
         gc.collect()
-
+        
         try:
             torch.cuda.reset_peak_memory_stats(device)
         except:
diff --git a/wanvideo/modules/attention.py b/wanvideo/modules/attention.py
index dde4cce..af2aa75 100644
--- a/wanvideo/modules/attention.py
+++ b/wanvideo/modules/attention.py
@@ -173,10 +173,10 @@ def attention(
             version=fa_version,
         )
     elif attention_mode == 'sdpa':
-        if q_lens is not None or k_lens is not None:
-            warnings.warn(
-                'Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance.'
-            )
+        # if q_lens is not None or k_lens is not None:
+        #     warnings.warn(
+        #         'Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance.'
+        #     )
         attn_mask = None
 
         q = q.transpose(1, 2).to(dtype)
diff --git a/wanvideo/modules/model.py b/wanvideo/modules/model.py
index 329d31c..a063acb 100644
--- a/wanvideo/modules/model.py
+++ b/wanvideo/modules/model.py
@@ -273,7 +273,8 @@ class WanAttentionBlock(nn.Module):
                                                                       num_heads,
                                                                       (-1, -1),
                                                                       qk_norm,
-                                                                      eps)
+                                                                      eps,#attention_mode=attention_mode sageattn doesn't seem faster here
+                                                                      )
         self.norm2 = WanLayerNorm(dim, eps)
         self.ffn = nn.Sequential(
             nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'),