remove unused function.

Author: sayakpaul

src/diffusers/models/attention_dispatch.py

@@ -589,25 +589,6 @@ def _normalize_attn_mask(attn_mask: torch.Tensor, batch_size: int, seq_len_k: in
     return attn_mask
 
 
-def _maybe_unflatten_attention_heads(out: torch.Tensor, reference_q: torch.Tensor) -> torch.Tensor:
-    """
-    Flash Attention 3 (and some hub builds) may return tensors where the head and head-dim axes are packed together.
-    Use the original query to restore the canonical [B, S, H, D] shape expected by the rest of the codebase.
-    """
-    if reference_q.ndim != 4 or out.ndim != 3:
-        return out
-
-    if out.shape[0] != reference_q.shape[0] or out.shape[1] != reference_q.shape[1]:
-        return out
-
-    num_heads, head_dim = reference_q.shape[-2:]
-    expected_width = num_heads * head_dim
-    if out.shape[-1] != expected_width:
-        return out
-
-    return out.reshape(reference_q.shape[0], reference_q.shape[1], num_heads, head_dim)
-
-
 def _flex_attention_causal_mask_mod(batch_idx, head_idx, q_idx, kv_idx):
     return q_idx >= kv_idx