Update

Author: ParagEkbote

src/diffusers/models/attention_dispatch.py

@@ -576,6 +576,9 @@ def _flash_attention(
     query: torch.Tensor,
     key: torch.Tensor,
     value: torch.Tensor,
+    dropout_p: float = 0.0,
+    scale: Optional[float] = None,
+    is_causal: bool = False,
     window_size: Tuple[int, int] = (-1, -1),
     softcap: float = 0.0,
     alibi_slopes: Optional[torch.Tensor] = None,
@@ -586,6 +589,9 @@ def _flash_attention(
         q=query,
         k=key,
         v=value,
+        dropout_p=dropout_p,
+        softmax_scale=scale,
+        causal=is_causal,
         window_size=window_size,
         softcap=softcap,
         alibi_slopes=alibi_slopes,
@@ -748,7 +754,6 @@ def _flash_attention_hub(
     query: torch.Tensor,
     key: torch.Tensor,
     value: torch.Tensor,
-    dropout_p: float = 0.0,
     scale: Optional[float] = None,
     is_causal: bool = False,
     window_size: Tuple[int, int] = (-1, -1),
@@ -761,7 +766,6 @@ def _flash_attention_hub(
         q=query,
         k=key,
         v=value,
-        dropout_p=dropout_p,
         softmax_scale=scale,
         causal=is_causal,
         window_size=window_size,