From 047775615bdd465e7d5059082ab2c0db46caab63 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 8 Feb 2023 14:24:27 -0500 Subject: [PATCH] Lower the chances of an OOM. --- comfy/ldm/modules/sub_quadratic_attention.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/comfy/ldm/modules/sub_quadratic_attention.py b/comfy/ldm/modules/sub_quadratic_attention.py index c0ee0cf2..45396d86 100644 --- a/comfy/ldm/modules/sub_quadratic_attention.py +++ b/comfy/ldm/modules/sub_quadratic_attention.py @@ -76,7 +76,8 @@ def _summarize_chunk( ) max_score, _ = torch.max(attn_weights, -1, keepdim=True) max_score = max_score.detach() - exp_weights = torch.exp(attn_weights - max_score) + torch.exp(attn_weights - max_score, out=attn_weights) + exp_weights = attn_weights exp_values = torch.bmm(exp_weights, value) max_score = max_score.squeeze(-1) return AttnChunk(exp_values, exp_weights.sum(dim=-1), max_score)