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

Try again with vae tiled decoding if regular fails because of OOM.

Browse Source
This commit is contained in:
comfyanonymous
2023-03-22 14:49:00 -04:00
parent aae9fe0cf9
commit 3ed4a4e4e6
5 changed files with 28 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -24,10 +24,7 @@ except ImportError:
from torch import Tensor
from typing import List
try:
OOM_EXCEPTION = torch.cuda.OutOfMemoryError
except:
OOM_EXCEPTION = Exception
import model_management
def dynamic_slice(
x: Tensor,
@@ -161,7 +158,7 @@ def _get_attention_scores_no_kv_chunking(
try:
attn_probs = attn_scores.softmax(dim=-1)
del attn_scores
except OOM_EXCEPTION:
except model_management.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)