298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-01-10 18:05:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

This makes pytorch2.0 attention perform a bit faster.

Browse Source
This commit is contained in:
comfyanonymous 2023-04-22 14:30:39 -04:00
parent 989acd769a
commit 6908f9c949
1 changed files with 2 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
comfy/ldm/modules/attention.py
View File

@ -455,11 +455,7 @@ class CrossAttentionPytorch(nn.Module):
b, _, _ = q.shape
q, k, v = map(
lambda t: t.unsqueeze(3)
.reshape(b, t.shape[1], self.heads, self.dim_head)
.permute(0, 2, 1, 3)
.reshape(b * self.heads, t.shape[1], self.dim_head)
.contiguous(),
lambda t: t.view(b, -1, self.heads, self.dim_head).transpose(1, 2),
(q, k, v),
)
@ -468,10 +464,7 @@ class CrossAttentionPytorch(nn.Module):
if exists(mask):
raise NotImplementedError
out = (
out.unsqueeze(0)
.reshape(b, self.heads, out.shape[1], self.dim_head)
.permute(0, 2, 1, 3)
.reshape(b, out.shape[1], self.heads * self.dim_head)
out.transpose(1, 2).reshape(b, -1, self.heads * self.dim_head)
)
return self.to_out(out)