Empty cache on mps.

2 years ago · 94680732d3
parent d200fa1314
commit 94680732d3
1 changed files with 4 additions and 1 deletions
--- a/comfy/model_management.py
+++ b/comfy/model_management.py
@ -389,7 +389,10 @@ def should_use_fp16():

 def soft_empty_cache():
    global xpu_available
-    if xpu_available:
+    global vram_state
+    if vram_state == VRAMState.MPS:
+        torch.mps.empty_cache()
+    elif xpu_available:
        torch.xpu.empty_cache()
    elif torch.cuda.is_available():
        if torch.version.cuda: #This seems to make things worse on ROCm so I only do it for cuda