Update modules/model.py

modules/model.py

@@ -203,14 +203,14 @@ class CrossAttnProcessor(nn.Module):
             k_bucket_size = 1024
 
             # use flash-attention
-            hidden_states = FlashAttn.apply(
+            hidden_states = FlashAttentionFunction.apply(
                 query.contiguous(),
                 key.contiguous(),
                 value.contiguous(),
                 attention_mask,
-                causal=False,
-                q_bucket_size=q_bucket_size,
-                k_bucket_size=k_bucket_size,
+                False,
+                q_bucket_size,
+                k_bucket_size,
             )
             hidden_states = hidden_states.to(query.dtype)
 
@@ -1021,5 +1021,3 @@ class FlashAttentionFunction(Function):
                 dvc.add_(dv_chunk)
 
         return dq, dk, dv, None, None, None, None
-    
-FlashAttn = FlashAttentionFunction()