xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-09-13 17:07:23 +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

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

@ -88,7 +88,8 @@ def Normalize(in_channels, dtype=None, device=None):
def attention_basic(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
dim_head //= heads
scale = dim_head ** -0.5
@ -103,11 +104,8 @@ def attention_basic(q, k, v, heads, mask=None, attn_precision=None):
(q, k, v),
)
# force cast to fp32 to avoid overflowing
if attn_precision == torch.float32:
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
# force cast to fp32 to avoid overflowing if args.dont_upcast_attention is not set
sim = einsum('b i d, b j d -> b i j', q.to(dtype=cast_to_type), k.to(dtype=cast_to_type) * scale
del q, k
@ -262,7 +260,7 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None):
end = i + slice_size
if upcast:
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:
s1 = einsum('b i d, b j d -> b i j', q[:, i:end], k) * scale
@ -312,6 +310,9 @@ except:
pass
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
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)
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),
)
@ -347,10 +348,13 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None):
return out
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
dim_head //= heads
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),
)