From da8916f92649fc4d947cb46d9d8f8ea1621b2a59 Mon Sep 17 00:00:00 2001
From: AUTOMATIC1111 <16777216c@gmail.com>
Date: Sat, 8 Jul 2023 17:13:18 +0300
Subject: added torch.mps.empty_cache() to torch_gc() changed a bunch of places
 that use torch.cuda.empty_cache() to use torch_gc() instead

---
 modules/codeformer_model.py | 2 +-
 modules/devices.py          | 3 +++
 modules/sd_models.py        | 1 -
 3 files changed, 4 insertions(+), 2 deletions(-)

(limited to 'modules')

diff --git a/modules/codeformer_model.py b/modules/codeformer_model.py
index f293acf5..da42b5e9 100644
--- a/modules/codeformer_model.py
+++ b/modules/codeformer_model.py
@@ -99,7 +99,7 @@ def setup_model(dirname):
                             output = self.net(cropped_face_t, w=w if w is not None else shared.opts.code_former_weight, adain=True)[0]
                             restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
                         del output
-                        torch.cuda.empty_cache()
+                        devices.torch_gc()
                     except Exception:
                         errors.report('Failed inference for CodeFormer', exc_info=True)
                         restored_face = tensor2img(cropped_face_t, rgb2bgr=True, min_max=(-1, 1))
diff --git a/modules/devices.py b/modules/devices.py
index 620ed1a6..c5ad950f 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -49,10 +49,13 @@ def get_device_for(task):
 
 
 def torch_gc():
+
     if torch.cuda.is_available():
         with torch.cuda.device(get_cuda_device_string()):
             torch.cuda.empty_cache()
             torch.cuda.ipc_collect()
+    elif has_mps() and hasattr(torch.mps, 'empty_cache'):
+        torch.mps.empty_cache()
 
 
 def enable_tf32():
diff --git a/modules/sd_models.py b/modules/sd_models.py
index f65f4e36..653c4cc0 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -590,7 +590,6 @@ def unload_model_weights(sd_model=None, info=None):
         sd_model = None
         gc.collect()
         devices.torch_gc()
-        torch.cuda.empty_cache()
 
     print(f"Unloaded weights {timer.summary()}.")
 
-- 
cgit v1.2.1