From 3c70da41928e358a217163f56789e0c3265cfd9c Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Wed, 26 Feb 2025 12:17:52 +0200
Subject: [PATCH] Update attention.py

---
 wanvideo/modules/attention.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/wanvideo/modules/attention.py b/wanvideo/modules/attention.py
index 1cc45ee..dde4cce 100644
--- a/wanvideo/modules/attention.py
+++ b/wanvideo/modules/attention.py
@@ -153,9 +153,9 @@ def attention(
     attention_mode='sdpa',
 ):  
     if "flash" in attention_mode:
-        if attention_mode == 'flash_attention_2':
+        if attention_mode == 'flash_attn_2':
             fa_version = 2
-        elif attention_mode == 'flash_attention_3':
+        elif attention_mode == 'flash_attn_3':
             fa_version = 3
         return flash_attention(
             q=q,