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

Implement attention mask on xformers.

Browse Source
This commit is contained in:
comfyanonymous 2024-01-06 04:33:03 -05:00
parent af94eb14e3
commit 3ad0191bfb
1 changed files with 7 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -294,11 +294,14 @@ def attention_xformers(q, k, v, heads, mask=None):
(q, k, v),
)
# actually compute the attention, what we cannot get enough of
out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None)
if mask is not None:
pad = 8 - q.shape[1] % 8
mask_out = torch.empty([q.shape[0], q.shape[1], q.shape[1] + pad], dtype=q.dtype, device=q.device)
mask_out[:, :, :mask.shape[-1]] = mask
mask = mask_out[:, :, :mask.shape[-1]]
out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=mask)
if exists(mask):
raise NotImplementedError
out = (
out.unsqueeze(0)
.reshape(b, heads, -1, dim_head)