fa

Author: zpcore

src/diffusers/models/attention_processor.py

@@ -24,6 +24,7 @@
 from ..utils.import_utils import is_torch_npu_available, is_xformers_available
 from ..utils.torch_utils import is_torch_version, maybe_allow_in_graph
 
+from torch_xla.experimental.custom_kernel import flash_attention
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
@@ -2364,9 +2365,17 @@ def __call__(
 
         # the output of sdp = (batch, num_heads, seq_len, head_dim)
         # TODO: add support for attn.scale when we move to Torch 2.1
-        hidden_states = F.scaled_dot_product_attention(
-            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
-        )
+       
+        # print("piz_test6")
+        # logger.warning(f"piz_debug2: query shape: {query.shape}, key shape: {key.shape}, value shape: {value.shape}")
+        if attention_mask is not None:
+            attention_mask = attention_mask.unsqueeze(1).expand(-1, attn.heads, -1, -1)
+            hidden_states = hidden_states * attention_mask
+        
+        hidden_states = flash_attention(query, key, value, causal=False)
+        # hidden_states = F.scaled_dot_product_attention(
+        #     query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        # )
 
         hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
         hidden_states = hidden_states.to(query.dtype)