298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-03-13 14:21:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

Make xformers work with hypertile.

Browse Source
This commit is contained in:
comfyanonymous 2023-10-21 13:23:03 -04:00
parent 1443caf373
commit 9906e3efe3
1 changed files with 7 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -253,12 +253,14 @@ def attention_split(q, k, v, heads, mask=None):
return r2
def attention_xformers(q, k, v, heads, mask=None):
b, _, _ = q.shape
b, _, dim_head = q.shape
dim_head //= heads
q, k, v = map(
lambda t: t.unsqueeze(3)
.reshape(b, t.shape[1], heads, -1)
.reshape(b, -1, heads, dim_head)
.permute(0, 2, 1, 3)
.reshape(b * heads, t.shape[1], -1)
.reshape(b * heads, -1, dim_head)
.contiguous(),
(q, k, v),
)
@ -270,9 +272,9 @@ def attention_xformers(q, k, v, heads, mask=None):
raise NotImplementedError
out = (
out.unsqueeze(0)
.reshape(b, heads, out.shape[1], -1)
.reshape(b, heads, -1, dim_head)
.permute(0, 2, 1, 3)
.reshape(b, out.shape[1], -1)
.reshape(b, -1, heads * dim_head)
)
return out