Update attention_processor.py

Author: hlky

src/diffusers/models/attention_processor.py

@@ -612,7 +612,7 @@ def batch_to_head_dim(self, tensor: torch.Tensor) -> torch.Tensor:
 
     def head_to_batch_dim(self, tensor: torch.Tensor, out_dim: int = 3) -> torch.Tensor:
         r"""
-        Reshape the tensor from `[batch_size, seq_len, dim]` to 
+        Reshape the tensor from `[batch_size, seq_len, dim]` to
         `[batch_size, seq_len, heads, dim // heads]`  for out_dim==4
         or  `[batch_size * heads, seq_len, dim // heads]` for out_dim==3
         where `heads` is the number of heads initialized while constructing the `Attention` class.
@@ -625,15 +625,16 @@ def head_to_batch_dim(self, tensor: torch.Tensor, out_dim: int = 3) -> torch.Ten
         Returns:
             `torch.Tensor`: The reshaped tensor.
         """
+        if out_dim not in {3, 4}:
+            raise ValueError(f"Expected `out_dim` to be 3 or 4, got {out_dim}.")
         head_size = self.heads
         if tensor.ndim == 3:
             batch_size, seq_len, dim = tensor.shape
             extra_dim = 1
         else:
             batch_size, extra_dim, seq_len, dim = tensor.shape
         tensor = tensor.reshape(batch_size, seq_len * extra_dim, head_size, dim // head_size)
-        
-        assert out_dim in [3,4]
+
         if out_dim == 3:
             tensor = tensor.permute(0, 2, 1, 3)
             tensor = tensor.reshape(batch_size * head_size, seq_len * extra_dim, dim // head_size)