298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-04-19 02:43:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

Lower the chances of an OOM.

Browse Source
This commit is contained in:
comfyanonymous 2023-02-08 14:24:27 -05:00
parent 853e96ada3
commit 047775615b
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
comfy/ldm/modules/sub_quadratic_attention.py
View File

@ -76,7 +76,8 @@ def _summarize_chunk(
) )
max_score, _ = torch.max(attn_weights, -1, keepdim=True) max_score, _ = torch.max(attn_weights, -1, keepdim=True)
max_score = max_score.detach() 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) exp_values = torch.bmm(exp_weights, value)
max_score = max_score.squeeze(-1) max_score = max_score.squeeze(-1)
return AttnChunk(exp_values, exp_weights.sum(dim=-1), max_score) return AttnChunk(exp_values, exp_weights.sum(dim=-1), max_score)