From 1a4edd19cd2cc82f1d7e0796181a826299f7304e Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Fri, 10 Feb 2023 11:47:41 -0500 Subject: [PATCH] Fix overflow issue with inplace softmax. --- comfy/ldm/modules/sub_quadratic_attention.py | 1 + 1 file changed, 1 insertion(+) diff --git a/comfy/ldm/modules/sub_quadratic_attention.py b/comfy/ldm/modules/sub_quadratic_attention.py index 6f5113383..5abed48c6 100644 --- a/comfy/ldm/modules/sub_quadratic_attention.py +++ b/comfy/ldm/modules/sub_quadratic_attention.py @@ -158,6 +158,7 @@ def _get_attention_scores_no_kv_chunking( del attn_scores except OOM_EXCEPTION: print("ran out of memory while running softmax in _get_attention_scores_no_kv_chunking, trying slower in place softmax instead") + attn_scores -= attn_scores.max(dim=-1, keepdim=True).values torch.exp(attn_scores, out=attn_scores) summed = torch.sum(attn_scores, dim=-1, keepdim=True) attn_scores /= summed