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] 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),
     )