jinx/ComfyUI
1
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-08-02 23:14:49 +08:00
Code Issues Activity

Fix some OOM issues with split and sub quad attention.

Browse Source
This commit is contained in:
comfyanonymous
2023-10-25 20:17:28 -04:00
parent 7fbb217d3a
commit a373367b0c
2 changed files with 9 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -83,7 +83,8 @@ def _summarize_chunk(
)
max_score, _ = torch.max(attn_weights, -1, keepdim=True)
max_score = max_score.detach()
torch.exp(attn_weights - max_score, out=attn_weights)
attn_weights -= max_score
torch.exp(attn_weights, out=attn_weights)
exp_weights = attn_weights.to(value.dtype)
exp_values = torch.bmm(exp_weights, value)
max_score = max_score.squeeze(-1)