From 8299ebdaaee68966b10e40e3e5e27caeb4117d0f Mon Sep 17 00:00:00 2001 From: shawnington <88048838+shawnington@users.noreply.github.com> Date: Mon, 3 Jun 2024 08:21:35 -0700 Subject: [PATCH 1/4] Update attention.py to respect FORCE_UPCAST_ATTENTION_DTYPE 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` --- comfy/ldm/modules/attention.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index da9f7aab7..d03944c99 100644 --- a/comfy/ldm/modules/attention.py +++ b/comfy/ldm/modules/attention.py @@ -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), ) From 938f824aae43816ef5d3e0ef15ed1d0b3be57268 Mon Sep 17 00:00:00 2001 From: shawnington <88048838+shawnington@users.noreply.github.com> Date: Mon, 3 Jun 2024 12:23:47 -0700 Subject: [PATCH 2/4] fixed typo in function --- comfy/ldm/modules/attention.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index d03944c99..363dd89bb 100644 --- a/comfy/ldm/modules/attention.py +++ b/comfy/ldm/modules/attention.py @@ -105,8 +105,8 @@ def attention_basic(q, k, v, heads, mask=None, attn_precision=None): ) # 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 - + 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 if exists(mask): From 438b1ea3990d775bfa99cb04a834b051dc7599f6 Mon Sep 17 00:00:00 2001 From: shawnington <88048838+shawnington@users.noreply.github.com> Date: Mon, 3 Jun 2024 12:57:46 -0700 Subject: [PATCH 3/4] fixed missing closing bracket --- comfy/ldm/modules/attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index 363dd89bb..b321bb081 100644 --- a/comfy/ldm/modules/attention.py +++ b/comfy/ldm/modules/attention.py @@ -260,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].to(dtype=torch.float32), k.to(dtype=torch.float32) * 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 From efd589391347cfb084c4db80a495da2c7e288d53 Mon Sep 17 00:00:00 2001 From: shawnington <88048838+shawnington@users.noreply.github.com> Date: Sun, 16 Jun 2024 11:43:41 -0400 Subject: [PATCH 4/4] Update attention.py to resolve conflicts with main branch updates --- comfy/ldm/modules/attention.py | 57 ++++++++++++++++++++++------------ 1 file changed, 38 insertions(+), 19 deletions(-) diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index b321bb081..e1e918273 100644 --- a/comfy/ldm/modules/attention.py +++ b/comfy/ldm/modules/attention.py @@ -86,12 +86,16 @@ class FeedForward(nn.Module): 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) -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) cast_to_type = attn_precision if attn_precision is not None else q.dtype - - b, _, dim_head = q.shape - dim_head //= heads + + if skip_reshape: + b, _, _, dim_head = q.shape + else: + b, _, dim_head = q.shape + dim_head //= heads + scale = dim_head ** -0.5 h = heads @@ -309,12 +313,15 @@ try: except: 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) cast_to_type = attn_precision if attn_precision is not None else q.dtype - b, _, dim_head = q.shape - dim_head //= heads + if skip_reshape: + b, _, _, dim_head = q.shape + else: + b, _, dim_head = q.shape + dim_head //= heads disabled_xformers = False @@ -329,10 +336,16 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None): if disabled_xformers: return attention_pytorch(q, k, v, heads, mask) - q, k, v = map( - lambda t: t.reshape(b, -1, heads, dim_head).to(dtype=cast_to_type), - (q, k, v), - ) + 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( + lambda t: t.reshape(b, -1, heads, dim_head).to(dtype=cast_to_type), + (q, k, v), + ) if mask is not None: pad = 8 - q.shape[1] % 8 @@ -347,16 +360,22 @@ 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): +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 - - b, _, dim_head = q.shape - dim_head //= heads - q, k, v = map( - lambda t: t.view(b, -1, heads, dim_head).transpose(1, 2).to(dtype=cast_to_type), - (q, k, v), - ) + + 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 + dim_head //= heads + q, k, v = map( + lambda t: t.view(b, -1, heads, dim_head).transpose(1, 2).to(dtype=cast_to_type), + (q, k, v), + ) out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False) out = (