From 73fdf050f92e0088346c0a87252a53958cee313d Mon Sep 17 00:00:00 2001
From: ControlAltAI <6172927+gseth@users.noreply.github.com>
Date: Sat, 21 Dec 2024 02:45:49 +0530
Subject: [PATCH] Flux Attention Control updated to work with Latest version of
 ComfyUI

---
 flux_attention_control_node.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/flux_attention_control_node.py b/flux_attention_control_node.py
index 13c667c..de24111 100644
--- a/flux_attention_control_node.py
+++ b/flux_attention_control_node.py
@@ -187,15 +187,21 @@ class FluxAttentionControl:
         return cross_mask, q_scale, k_scale
 
     def xformers_attention(self, q: Tensor, k: Tensor, v: Tensor, pe: Tensor,
-                          attn_mask: Optional[Tensor] = None) -> Tensor:
+                            attn_mask: Optional[Tensor] = None,
+                            mask: Optional[Tensor] = None) -> Tensor:  # Added mask parameter
         q, k = flux_math.apply_rope(q, k, pe)
         q = rearrange(q, "B H L D -> B L H D")
         k = rearrange(k, "B H L D -> B L H D")
         v = rearrange(v, "B H L D -> B L H D")
-        if attn_mask is not None:
-            x = xattention(q, k, v, attn_bias=attn_mask)
+        
+        # Use attn_mask if provided, otherwise use the mask parameter
+        attention_bias = attn_mask if attn_mask is not None else mask
+        
+        if attention_bias is not None:
+            x = xattention(q, k, v, attn_bias=attention_bias)
         else:
             x = xattention(q, k, v)
+            
         x = rearrange(x, "B L H D -> B L (H D)")
         return x