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

Update attention.py to respect FORCE_UPCAST_ATTENTION_DTYPE

Browse Source 
Fixed attention precision not being cast in attention_pytorch and others functions. This led to the functions not being able to respect the `--dont-upcast-attention` flag.  

change calls to `.float()` to `.to(dtype=torch.float32)` in several locations, as it profiles much faster. 

removed unneeded check for `attn_precision == torch.float32`, as the change from `.float()` to `.to(dtype=cast_to_type)` does not cast or copy if `tensor.dtype == cast_to_type`
This commit is contained in:
shawnington 2024-06-03 08:21:35 -07:00 committed by GitHub
parent 809cc85a8e
commit 8299ebdaae
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 13 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
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): def attention_basic(q, k, v, heads, mask=None, attn_precision=None):
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
b, _, dim_head = q.shape b, _, dim_head = q.shape
dim_head //= heads dim_head //= heads
@ -103,11 +104,8 @@ def attention_basic(q, k, v, heads, mask=None, attn_precision=None):
(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
@ -262,7 +260,7 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None):
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
@ -312,6 +310,9 @@ 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):
attn_precision = get_attn_precision(attn_precision)
cast_to_type = attn_precision if attn_precision is not None else q.dtype
b, _, dim_head = q.shape b, _, dim_head = q.shape
dim_head //= heads dim_head //= heads
@ -329,7 +330,7 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None):
return attention_pytorch(q, k, v, heads, mask) return attention_pytorch(q, k, v, heads, mask)
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),
) )
@ -347,10 +348,13 @@ 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):
attn_precision = get_attn_precision(attn_precision)
cast_to_type = attn_precision if attn_precision is not None else q.dtype
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),
) )