From 0a020a6f1848c6aba0be7e7c4e8ecb727a79d920 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Wed, 26 Feb 2025 12:15:35 +0200
Subject: [PATCH] maybe fix flash attention

---
 wanvideo/modules/attention.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/wanvideo/modules/attention.py b/wanvideo/modules/attention.py
index 0735663..1cc45ee 100644
--- a/wanvideo/modules/attention.py
+++ b/wanvideo/modules/attention.py
@@ -151,11 +151,12 @@ def attention(
     deterministic=False,
     dtype=torch.bfloat16,
     attention_mode='sdpa',
-):
-    if attention_mode == 'flash_attention_2':
-        fa_version = 2
-    elif attention_mode == 'flash_attention_3':
-        fa_version = 3
+):  
+    if "flash" in attention_mode:
+        if attention_mode == 'flash_attention_2':
+            fa_version = 2
+        elif attention_mode == 'flash_attention_3':
+            fa_version = 3
         return flash_attention(
             q=q,
             k=k,