From 24bb774432808de43d06dc0d4cdd740ab5e40bca Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Wed, 26 Feb 2025 09:37:50 +0200
Subject: [PATCH] Update model_optimization_nodes.py

---
 nodes/model_optimization_nodes.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/nodes/model_optimization_nodes.py b/nodes/model_optimization_nodes.py
index 1dc0fab..fabb9a3 100644
--- a/nodes/model_optimization_nodes.py
+++ b/nodes/model_optimization_nodes.py
@@ -83,6 +83,7 @@ class BaseLoaderKJ:
             comfy.ldm.flux.math.optimized_attention = attention_sage
             comfy.ldm.genmo.joint_model.asymm_models_joint.optimized_attention = attention_sage
             comfy.ldm.cosmos.blocks.optimized_attention = attention_sage
+            comfy.ldm.wan.model.optimized_attention = attention_sage
 
         else:
             comfy_attention.optimized_attention = orig_attention
@@ -90,6 +91,7 @@ class BaseLoaderKJ:
             comfy.ldm.flux.math.optimized_attention = orig_attention
             comfy.ldm.genmo.joint_model.asymm_models_joint.optimized_attention = orig_attention
             comfy.ldm.cosmos.blocks.optimized_attention = orig_attention
+            comfy.ldm.wan.model.optimized_attention = orig_attention
 
         if patch_cublaslinear:
             if not BaseLoaderKJ.cublas_patched: