Merge pull request #1616 from brkirch/cpu-cmdline-opt

Add --use-cpu command line option

1 files changed, 4 insertions, 5 deletions

diff --git a/modules/esrgan_model.py b/modules/esrgan_model.py
index 4aed9283..d17e730f 100644
--- a/modules/esrgan_model.py
+++ b/modules/esrgan_model.py
@@ -6,8 +6,7 @@ from PIL import Image
 from basicsr.utils.download_util import load_file_from_url

 

 import modules.esrgam_model_arch as arch

-from modules import shared, modelloader, images

-from modules.devices import has_mps

+from modules import shared, modelloader, images, devices

 from modules.paths import models_path

 from modules.upscaler import Upscaler, UpscalerData

 from modules.shared import opts

@@ -97,7 +96,7 @@ class UpscalerESRGAN(Upscaler):
         model = self.load_model(selected_model)

         if model is None:

             return img

-        model.to(shared.device)

+        model.to(devices.device_esrgan)

         img = esrgan_upscale(model, img)

         return img

 

@@ -112,7 +111,7 @@ class UpscalerESRGAN(Upscaler):
             print("Unable to load %s from %s" % (self.model_path, filename))

             return None

 

-        pretrained_net = torch.load(filename, map_location='cpu' if has_mps else None)

+        pretrained_net = torch.load(filename, map_location='cpu' if shared.device.type == 'mps' else None)

         crt_model = arch.RRDBNet(3, 3, 64, 23, gc=32)

 

         pretrained_net = fix_model_layers(crt_model, pretrained_net)

@@ -127,7 +126,7 @@ def upscale_without_tiling(model, img):
     img = img[:, :, ::-1]

     img = np.moveaxis(img, 2, 0) / 255

     img = torch.from_numpy(img).float()

-    img = img.unsqueeze(0).to(shared.device)

+    img = img.unsqueeze(0).to(devices.device_esrgan)

     with torch.no_grad():

         output = model(img)

     output = output.squeeze().float().cpu().clamp_(0, 1).numpy()