xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-10-11 11:17:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Update attention.py to resolve conflicts with main branch updates

Browse Source
This commit is contained in:
shawnington 2024-06-16 11:43:41 -04:00 committed by GitHub
parent 438b1ea399
commit efd5893913
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 38 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -86,12 +86,16 @@ class FeedForward(nn.Module):
def Normalize(in_channels, dtype=None, device=None): def Normalize(in_channels, dtype=None, device=None):
return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True, dtype=dtype, device=device) return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True, dtype=dtype, device=device)
def attention_basic(q, k, v, heads, mask=None, attn_precision=None): def attention_basic(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False):
attn_precision = get_attn_precision(attn_precision) attn_precision = get_attn_precision(attn_precision)
cast_to_type = attn_precision if attn_precision is not None else q.dtype cast_to_type = attn_precision if attn_precision is not None else q.dtype
if skip_reshape:
b, _, _, dim_head = q.shape
else:
b, _, dim_head = q.shape b, _, dim_head = q.shape
dim_head //= heads dim_head //= heads
scale = dim_head ** -0.5 scale = dim_head ** -0.5
h = heads h = heads
@ -309,10 +313,13 @@ try:
except: except:
pass pass
def attention_xformers(q, k, v, heads, mask=None, attn_precision=None): def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False):
attn_precision = get_attn_precision(attn_precision) attn_precision = get_attn_precision(attn_precision)
cast_to_type = attn_precision if attn_precision is not None else q.dtype cast_to_type = attn_precision if attn_precision is not None else q.dtype
if skip_reshape:
b, _, _, dim_head = q.shape
else:
b, _, dim_head = q.shape b, _, dim_head = q.shape
dim_head //= heads dim_head //= heads
@ -329,6 +336,12 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None):
if disabled_xformers: if disabled_xformers:
return attention_pytorch(q, k, v, heads, mask) return attention_pytorch(q, k, v, heads, mask)
if skip_reshape:
q, k, v = map(
lambda t: t.reshape(b * heads, -1, dim_head).to(dtype=cast_to_type),
(q, k, v),
)
else:
q, k, v = map( q, k, v = map(
lambda t: t.reshape(b, -1, heads, dim_head).to(dtype=cast_to_type), lambda t: t.reshape(b, -1, heads, dim_head).to(dtype=cast_to_type),
(q, k, v), (q, k, v),
@ -347,10 +360,16 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None):
) )
return out return out
def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None): def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False):
attn_precision = get_attn_precision(attn_precision) attn_precision = get_attn_precision(attn_precision)
cast_to_type = attn_precision if attn_precision is not None else q.dtype cast_to_type = attn_precision if attn_precision is not None else q.dtype
if skip_reshape:
b, _, _, dim_head = q.shape
q, k, v = map(
lambda t: t.to(dtype=cast_to_type), (q, k, v),
)
else:
b, _, dim_head = q.shape b, _, dim_head = q.shape
dim_head //= heads dim_head //= heads
q, k, v = map( q, k, v = map(