From 94680732d32b4b540251c122aee36df8d37266e1 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Thu, 1 Jun 2023 03:52:51 -0400
Subject: [PATCH] Empty cache on mps.

---
 comfy/model_management.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/comfy/model_management.py b/comfy/model_management.py
index 10a7067..60bcd78 100644
--- 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