From bf67a5dcf44c3dbd88d1913478d4e02477915f33 Mon Sep 17 00:00:00 2001
From: Aarni Koskela <akx@iki.fi>
Date: Mon, 29 May 2023 10:38:51 +0300
Subject: Upscaler.load_model: don't return None, just use exceptions

---
 modules/realesrgan_model.py | 33 +++++++++++++++------------------
 1 file changed, 15 insertions(+), 18 deletions(-)

(limited to 'modules/realesrgan_model.py')

diff --git a/modules/realesrgan_model.py b/modules/realesrgan_model.py
index 0d9c2e48..0700b853 100644
--- a/modules/realesrgan_model.py
+++ b/modules/realesrgan_model.py
@@ -9,7 +9,6 @@ from modules.shared import cmd_opts, opts
 from modules import modelloader, errors
 
 
-
 class UpscalerRealESRGAN(Upscaler):
     def __init__(self, path):
         self.name = "RealESRGAN"
@@ -43,9 +42,10 @@ class UpscalerRealESRGAN(Upscaler):
         if not self.enable:
             return img
 
-        info = self.load_model(path)
-        if not os.path.exists(info.local_data_path):
-            print(f"Unable to load RealESRGAN model: {info.name}")
+        try:
+            info = self.load_model(path)
+        except Exception:
+            errors.report(f"Unable to load RealESRGAN model {path}", exc_info=True)
             return img
 
         upsampler = RealESRGANer(
@@ -63,20 +63,17 @@ class UpscalerRealESRGAN(Upscaler):
         return image
 
     def load_model(self, path):
-        try:
-            info = next(iter([scaler for scaler in self.scalers if scaler.data_path == path]), None)
-
-            if info is None:
-                print(f"Unable to find model info: {path}")
-                return None
-
-            if info.local_data_path.startswith("http"):
-                info.local_data_path = modelloader.load_file_from_url(info.data_path, model_dir=self.model_download_path)
-
-            return info
-        except Exception:
-            errors.report("Error making Real-ESRGAN models list", exc_info=True)
-        return None
+        for scaler in self.scalers:
+            if scaler.data_path == path:
+                if scaler.local_data_path.startswith("http"):
+                    scaler.local_data_path = modelloader.load_file_from_url(
+                        scaler.data_path,
+                        model_dir=self.model_download_path,
+                    )
+                if not os.path.exists(scaler.local_data_path):
+                    raise FileNotFoundError(f"RealESRGAN data missing: {scaler.local_data_path}")
+                return scaler
+        raise ValueError(f"Unable to find model info: {path}")
 
     def load_models(self, _):
         return get_realesrgan_models(self)
-- 
cgit v1.2.1