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

Merge 01abd4b9b4ebab520dc64b606011edfdafc7fab3 into 8af9a91e0c47b9fc277077f2079873adf8edac05

Browse Source
This commit is contained in:
shawnington 2024-12-06 03:32:53 -08:00 committed by GitHub
commit e015462b06
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 19 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -88,6 +88,7 @@ def Normalize(in_channels, dtype=None, device=None):
def attention_basic(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False): 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
if skip_reshape: if skip_reshape:
b, _, _, dim_head = q.shape b, _, _, dim_head = q.shape
@ -113,12 +114,9 @@ def attention_basic(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
(q, k, v), (q, k, v),
) )
# force cast to fp32 to avoid overflowing # force cast to fp32 to avoid overflowing if args.dont_upcast_attention is not set
if attn_precision == torch.float32: sim = einsum('b i d, b j d -> b i j', q.to(dtype=cast_to_type), k.to(dtype=cast_to_type)) * scale
sim = einsum('b i d, b j d -> b i j', q.float(), k.float()) * scale
else:
sim = einsum('b i d, b j d -> b i j', q, k) * scale
del q, k del q, k
if exists(mask): if exists(mask):
@ -291,7 +289,7 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
end = i + slice_size end = i + slice_size
if upcast: if upcast:
with torch.autocast(enabled=False, device_type = 'cuda'): with torch.autocast(enabled=False, device_type = 'cuda'):
s1 = einsum('b i d, b j d -> b i j', q[:, i:end].float(), k.float()) * scale s1 = einsum('b i d, b j d -> b i j', q[:, i:end].to(dtype=torch.float32), k.to(dtype=torch.float32)) * scale
else: else:
s1 = einsum('b i d, b j d -> b i j', q[:, i:end], k) * scale s1 = einsum('b i d, b j d -> b i j', q[:, i:end], k) * scale
@ -344,6 +342,8 @@ except:
pass pass
def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False): def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False):
attn_precision = get_attn_precision(attn_precision)
cast_to_type = attn_precision if attn_precision is not None else q.dtype
if skip_reshape: if skip_reshape:
b, _, _, dim_head = q.shape b, _, _, dim_head = q.shape
else: else:
@ -365,12 +365,12 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_resh
if skip_reshape: if skip_reshape:
q, k, v = map( q, k, v = map(
lambda t: t.reshape(b * heads, -1, dim_head), lambda t: t.reshape(b * heads, -1, dim_head).to(dtype=cast_to_type),
(q, k, v), (q, k, v)
) )
else: else:
q, k, v = map( q, k, v = map(
lambda t: t.reshape(b, -1, heads, dim_head), lambda t: t.reshape(b, -1, heads, dim_head).to(dtype=cast_to_type),
(q, k, v), (q, k, v),
) )
@ -404,13 +404,20 @@ else:
def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False): def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False):
attn_precision = get_attn_precision(attn_precision)
cast_to_type = attn_precision if attn_precision is not None else q.dtype
if skip_reshape: if skip_reshape:
b, _, _, dim_head = q.shape b, _, _, dim_head = q.shape
q, k, v = map(
lambda t: t.to(dtype=cast_to_type), (q, k, v),
)
else: 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(
lambda t: t.view(b, -1, heads, dim_head).transpose(1, 2), lambda t: t.view(b, -1, heads, dim_head).transpose(1, 2).to(dtype=cast_to_type),
(q, k, v), (q, k, v),
) )