From b5d1af11b7dc718d4d91d379c75e46f4bd2e2fe6 Mon Sep 17 00:00:00 2001
From: Abdullah Barhoum <ka70911@gmail.com>
Date: Sun, 11 Sep 2022 07:11:27 +0200
Subject: [PATCH] Modular device management

---
 modules/devices.py      | 12 ++++++++++++
 modules/esrgan_model.py |  3 ++-
 modules/lowvram.py      |  8 ++------
 modules/shared.py       |  9 +++------
 4 files changed, 19 insertions(+), 13 deletions(-)
 create mode 100644 modules/devices.py

diff --git a/modules/devices.py b/modules/devices.py
new file mode 100644
index 00000000..25008a04
--- /dev/null
+++ b/modules/devices.py
@@ -0,0 +1,12 @@
+import torch
+
+
+# has_mps is only available in nightly pytorch (for now), `getattr` for compatibility
+has_mps = getattr(torch, 'has_mps', False)
+
+def get_optimal_device():
+  if torch.cuda.is_available():
+      return torch.device("cuda")
+  if has_mps:
+      return torch.device("mps")
+  return torch.device("cpu")
diff --git a/modules/esrgan_model.py b/modules/esrgan_model.py
index e86ad775..7f3baf31 100644
--- a/modules/esrgan_model.py
+++ b/modules/esrgan_model.py
@@ -9,12 +9,13 @@ from PIL import Image
 import modules.esrgam_model_arch as arch
 from modules import shared
 from modules.shared import opts
+from modules.devices import has_mps
 import modules.images
 
 
 def load_model(filename):
     # this code is adapted from https://github.com/xinntao/ESRGAN
-    pretrained_net = torch.load(filename, map_location='cpu' if torch.has_mps else None)
+    pretrained_net = torch.load(filename, map_location='cpu' if has_mps else None)
     crt_model = arch.RRDBNet(3, 3, 64, 23, gc=32)
 
     if 'conv_first.weight' in pretrained_net:
diff --git a/modules/lowvram.py b/modules/lowvram.py
index bd117491..079386c3 100644
--- a/modules/lowvram.py
+++ b/modules/lowvram.py
@@ -1,13 +1,9 @@
 import torch
+from modules.devices import get_optimal_device
 
 module_in_gpu = None
 cpu = torch.device("cpu")
-if torch.has_cuda:
-    device = gpu = torch.device("cuda")
-elif torch.has_mps:
-    device = gpu = torch.device("mps")
-else:
-    device = gpu = torch.device("cpu")
+device = gpu = get_optimal_device()
 
 def setup_for_low_vram(sd_model, use_medvram):
     parents = {}
diff --git a/modules/shared.py b/modules/shared.py
index 6ca9106c..74b0ad89 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -9,6 +9,7 @@ import tqdm
 
 import modules.artists
 from modules.paths import script_path, sd_path
+from modules.devices import get_optimal_device
 import modules.styles
 
 config_filename = "config.json"
@@ -43,12 +44,8 @@ parser.add_argument("--ui-config-file", type=str, help="filename to use for ui c
 
 cmd_opts = parser.parse_args()
 
-if torch.has_cuda:
-    device = torch.device("cuda")
-elif torch.has_mps:
-    device = torch.device("mps")
-else:
-    device = torch.device("cpu")
+device = get_optimal_device()
+
 batch_cond_uncond = cmd_opts.always_batch_cond_uncond or not (cmd_opts.lowvram or cmd_opts.medvram)
 parallel_processing_allowed = not cmd_opts.lowvram and not cmd_opts.medvram