From 00de4f5e0e2469d0e1c743b80e9be3a3349f1a2a Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Fri, 18 Apr 2025 11:47:19 +0300
Subject: [PATCH] Detect (a guess for now) TeaCache coefficients for SkyreelsV2

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

diff --git a/nodes.py b/nodes.py
index 034fc6d..2ca1395 100644
--- a/nodes.py
+++ b/nodes.py
@@ -586,17 +586,15 @@ class WanVideoModelLoader:
                 "e0": [8.10705460e+03, 2.13393892e+03, -3.72934672e+02, 1.66203073e+01, -4.17769401e-02],
             },
         }
-        
+        model_variant = "14B" #default to this
         if model_type == "i2v" or model_type == "fl2v":
-            if "480" in model or "fun" in model.lower() or "a2" in model.lower(): #just a guess for the Fun model for now...
+            if "480" in model or "fun" in model.lower() or "a2" in model.lower() or "540" in model: #just a guess for the Fun model for now...
                 model_variant = "i2v_480"
             elif "720" in model:
                 model_variant = "i2v_720"
         elif model_type == "t2v":
             model_variant = "14B"
-        else:
-            model_variant = "14B" #default to this
-            log.warning("Model variant not detected, defaulting TeaCache coefficients to 14B")
+            
         if dim == 1536:
             model_variant = "1_3B"
         log.info(f"Model variant detected: {model_variant}")