298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-01-25 15:55:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

torch.cuda.OutOfMemoryError is not present on older pytorch versions.

Browse Source
This commit is contained in:
comfyanonymous 2023-02-09 12:33:27 -05:00
parent 1d9ec62cfb
commit df40d4f3bf
1 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -19,6 +19,11 @@ from typing import Optional, NamedTuple, Protocol, List
from torch import Tensor from torch import Tensor
from typing import List from typing import List
try:
OOM_EXCEPTION = torch.cuda.OutOfMemoryError
except:
OOM_EXCEPTION = Exception
def dynamic_slice( def dynamic_slice(
x: Tensor, x: Tensor,
starts: List[int], starts: List[int],
@ -151,7 +156,7 @@ def _get_attention_scores_no_kv_chunking(
try: try:
attn_probs = attn_scores.softmax(dim=-1) attn_probs = attn_scores.softmax(dim=-1)
del attn_scores del attn_scores
except torch.cuda.OutOfMemoryError: except OOM_EXCEPTION:
print("ran out of memory while running softmax in _get_attention_scores_no_kv_chunking, trying slower in place softmax instead") print("ran out of memory while running softmax in _get_attention_scores_no_kv_chunking, trying slower in place softmax instead")
torch.exp(attn_scores, out=attn_scores) torch.exp(attn_scores, out=attn_scores)
summed = torch.sum(attn_scores, dim=-1, keepdim=True) summed = torch.sum(attn_scores, dim=-1, keepdim=True)