From b12b48e170ccff156dc6ec11242bb6af7d8437fd Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Thu, 14 Dec 2023 20:11:46 -0500 Subject: [PATCH 1/2] cleanup. --- comfy_extras/nodes_sag.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/comfy_extras/nodes_sag.py b/comfy_extras/nodes_sag.py index 4c609565ac4..0bcda84f08a 100644 --- a/comfy_extras/nodes_sag.py +++ b/comfy_extras/nodes_sag.py @@ -111,7 +111,6 @@ def patch(self, model, scale, blur_sigma): m = model.clone() attn_scores = None - mid_block_shape = None # TODO: make this work properly with chunked batches # currently, we can only save the attn from one UNet call @@ -134,7 +133,6 @@ def attn_and_record(q, k, v, extra_options): def post_cfg_function(args): nonlocal attn_scores - nonlocal mid_block_shape uncond_attn = attn_scores sag_scale = scale From a5056cfb1f41f1f9e6fcd523ef8091e6e7cd6e3b Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Fri, 15 Dec 2023 01:28:16 -0500 Subject: [PATCH 2/2] Remove useless code. --- comfy/ldm/modules/attention.py | 4 +--- comfy_extras/nodes_sag.py | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index 8d86aa53d2e..3e12886b07f 100644 --- a/comfy/ldm/modules/attention.py +++ b/comfy/ldm/modules/attention.py @@ -104,9 +104,7 @@ def attention_basic(q, k, v, heads, mask=None): # force cast to fp32 to avoid overflowing if _ATTN_PRECISION =="fp32": - with torch.autocast(enabled=False, device_type = 'cuda'): - q, k = q.float(), k.float() - sim = einsum('b i d, b j d -> b i j', q, k) * 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 diff --git a/comfy_extras/nodes_sag.py b/comfy_extras/nodes_sag.py index 0bcda84f08a..7e293ef63e4 100644 --- a/comfy_extras/nodes_sag.py +++ b/comfy_extras/nodes_sag.py @@ -27,9 +27,7 @@ def attention_basic_with_sim(q, k, v, heads, mask=None): # force cast to fp32 to avoid overflowing if _ATTN_PRECISION =="fp32": - with torch.autocast(enabled=False, device_type = 'cuda'): - q, k = q.float(), k.float() - sim = einsum('b i d, b j d -> b i j', q, k) * 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