for the (unused) attention bias which the baddbmm API requires us to pass in: create a smaller bias and broadcast it. this helps to prevent NaN result from baddbmm() on MPS on PyTorch 2.1.0.dev20230310.

Author: Birch-san

src/diffusers/models/cross_attention.py

@@ -231,16 +231,20 @@ def get_attention_scores(self, query, key, attention_mask=None):
             key = key.float()
 
         if attention_mask is None:
-            baddbmm_input = torch.empty(
-                query.shape[0], query.shape[1], key.shape[1], dtype=query.dtype, device=query.device
-            )
+            batch_x_heads, query_tokens, _ = query.shape
+            _, key_tokens, _ = key.shape
+            # expanding dims isn't strictly necessary (baddbmm supports broadcasting bias),
+            # but documents the expected shape without allocating any additional memory
+            attention_bias = torch.zeros(
+                1, 1, 1, dtype=query.dtype, device=query.device
+            ).expand(batch_x_heads, query_tokens, key_tokens)
             beta = 0
         else:
-            baddbmm_input = attention_mask
+            attention_bias = attention_mask
             beta = 1
 
         attention_scores = torch.baddbmm(
-            baddbmm_input,
+            attention_bias,
             query,
             key.transpose(-1, -2),
             beta=beta,