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

Add a force argument to soft_empty_cache to force a cache empty.

Browse Source
This commit is contained in:
comfyanonymous 2023-09-04 00:58:18 -04:00
parent 7746bdf7b0
commit 1938f5c5fe
3 changed files with 4 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
comfy/ldm/modules/attention.py
View File

@ -323,7 +323,7 @@ class CrossAttentionDoggettx(nn.Module):
break break
except model_management.OOM_EXCEPTION as e: except model_management.OOM_EXCEPTION as e:
if first_op_done == False: if first_op_done == False:
model_management.soft_empty_cache() model_management.soft_empty_cache(True)
if cleared_cache == False: if cleared_cache == False:
cleared_cache = True cleared_cache = True
print("out of memory error, emptying cache and trying again") print("out of memory error, emptying cache and trying again")

1
comfy/ldm/modules/diffusionmodules/model.py
View File

@ -186,6 +186,7 @@ def slice_attention(q, k, v):
del s2 del s2
break break
except model_management.OOM_EXCEPTION as e: except model_management.OOM_EXCEPTION as e:
model_management.soft_empty_cache(True)
steps *= 2 steps *= 2
if steps > 128: if steps > 128:
raise e raise e

4
comfy/model_management.py
View File

@ -639,14 +639,14 @@ def should_use_fp16(device=None, model_params=0, prioritize_performance=True):
return True return True
def soft_empty_cache(): def soft_empty_cache(force=False):
global cpu_state global cpu_state
if cpu_state == CPUState.MPS: if cpu_state == CPUState.MPS:
torch.mps.empty_cache() torch.mps.empty_cache()
elif is_intel_xpu(): elif is_intel_xpu():
torch.xpu.empty_cache() torch.xpu.empty_cache()
elif torch.cuda.is_available(): elif torch.cuda.is_available():
if is_nvidia(): #This seems to make things worse on ROCm so I only do it for cuda if force or is_nvidia(): #This seems to make things worse on ROCm so I only do it for cuda
torch.cuda.empty_cache() torch.cuda.empty_cache()
torch.cuda.ipc_collect() torch.cuda.ipc_collect()