kanshan/diffusers
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix CrossAttention._sliced_attention (#563)

Browse Source 
* Fix CrossAttention._sliced_attention

Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
This commit is contained in:
Yih-Dar 2022-09-19 18:07:32 +02:00 committed by GitHub
parent 8d36d5adb1
commit 84616b5de5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
src/diffusers/models/attention.py
View File

@ -249,13 +249,15 @@ class CrossAttention(nn.Module):
return tensor return tensor
def forward(self, hidden_states, context=None, mask=None): def forward(self, hidden_states, context=None, mask=None):
batch_size, sequence_length, dim = hidden_states.shape batch_size, sequence_length, _ = hidden_states.shape
query = self.to_q(hidden_states) query = self.to_q(hidden_states)
context = context if context is not None else hidden_states context = context if context is not None else hidden_states
key = self.to_k(context) key = self.to_k(context)
value = self.to_v(context) value = self.to_v(context)
dim = query.shape[-1]
query = self.reshape_heads_to_batch_dim(query) query = self.reshape_heads_to_batch_dim(query)
key = self.reshape_heads_to_batch_dim(key) key = self.reshape_heads_to_batch_dim(key)
value = self.reshape_heads_to_batch_dim(value) value = self.reshape_heads_to_batch_dim(value)