huggingface / huggingface/diffusers

【BUG】Attention.head_to_batch_dim has bug in terms of tensor permutation

Offen
#10,303 3 Kommentare 0 Reaktionen 0 zugewiesene Personen Auf GitHub ansehen
bug stale
Vorherrschende Sprache
Python
Sterne
34.5k
Forks
7.3k
Ø Merge
3 T. 3 Std.
Gemergte PRs (30 T.)
91

Beschreibung

### Describe the bug

https://github.com/huggingface/diffusers/blob/1826a1e7d31df48d345a20028b3ace48f09a4e60/src/diffusers/models/attention_processor.py#L613

when `out_dim==4`, the ourpout shape is mismatch to the function's comment ``[batch_size, seq_len, heads, dim // heads]`

here is original function
```
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 `[batch_size, seq_len, heads, dim // heads]` `heads` is
the number of heads initialized while constructing the `Attention` class.

Args:
tensor (`torch.Tensor`): The tensor to reshape.
out_dim (`int`, *optional*, defaults to `3`): The output dimension of the tensor. If `3`, the tensor is
reshaped to `[batch_size * heads, seq_len, dim // heads]`.

Returns:
`torch.Tensor`: The reshaped tensor.
"""
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)
tensor = tensor.permute(0, 2, 1, 3)

if out_dim == 3:
tensor = tensor.reshape(batch_size * head_size, seq_len * extra_dim, dim // head_size)

return tensor
```

and at [Line 633](https://github.com/huggingface/diffusers/blob/1826a1e7d31df48d345a20028b3ace48f09a4e60/src/diffusers/models/attention_processor.py#L633), `tensor = tensor.permute(0, 2, 1, 3)` the tensor permutes again

The correction should be moving Line633 to Line635.5 i.e.,
```
...
tensor = tensor.reshape(batch_size, seq_len * extra_dim, head_size, dim // head_size)

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)

return tensor
```

### Reproduction

just inside Attention, run
```
self.head_to_batch_dim(query,out_dim=4)
```

### Logs

_No response_

### System Info

This is irrelevant to the system & environment info

### Who can help?

_No response_

Beitragsleitfaden

Beitragsleitfaden öffnen

Rechercherichtung

Beginnen Sie in src/diffusers/models/attention_processor.py bei der verlinkten head_to_batch_dim-Implementierung und untersuchen Sie die Änderungen der Tensorform rund um die Permutation. Reproduzieren Sie das Problem, indem Sie Attention.head_to_batch_dim(query, out_dim=4) aufrufen, und überprüfen Sie anschließend, dass die zurückgegebene Form der dokumentierten Form entspricht, während das Verhalten von out_dim=3 weiterhin korrekt ist.

Vom Indexierungsmodell aus dem Issue-Text verfasst.

Bewertung

Tech-Stack
python, pytorch
Bereich
machine-learning
Issue-Typ
Bug
Schwierigkeit
3/5
Geschätzter Aufwand
1-2 Tage
Aktivitätsstatus
Veraltet
Klarheit
Klar beschrieben
Anfängerfreundlichkeit
45/100

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.