Merge branch 'dev' into master

1 files changed, 19 insertions, 26 deletions

diff --git a/modules/realesrgan_model.py b/modules/realesrgan_model.py
index 99983678..0700b853 100644
--- a/modules/realesrgan_model.py
+++ b/modules/realesrgan_model.py
@@ -1,15 +1,13 @@
 import os

-import sys

-import traceback

 

 import numpy as np

 from PIL import Image

-from basicsr.utils.download_util import load_file_from_url

 from realesrgan import RealESRGANer

 

 from modules.upscaler import Upscaler, UpscalerData

 from modules.shared import cmd_opts, opts

-from modules import modelloader

+from modules import modelloader, errors

+

 

 class UpscalerRealESRGAN(Upscaler):

     def __init__(self, path):

@@ -36,8 +34,7 @@ class UpscalerRealESRGAN(Upscaler):
                     self.scalers.append(scaler)

 

         except Exception:

-            print("Error importing Real-ESRGAN:", file=sys.stderr)

-            print(traceback.format_exc(), file=sys.stderr)

+            errors.report("Error importing Real-ESRGAN", exc_info=True)

             self.enable = False

             self.scalers = []

 

@@ -45,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(

@@ -65,21 +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 = load_file_from_url(url=info.data_path, model_dir=self.model_download_path, progress=True)

-

-            return info

-        except Exception as e:

-            print(f"Error making Real-ESRGAN models list: {e}", file=sys.stderr)

-            print(traceback.format_exc(), file=sys.stderr)

-        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)

@@ -135,5 +129,4 @@ def get_realesrgan_models(scaler):
         ]

         return models

     except Exception:

-        print("Error making Real-ESRGAN models list:", file=sys.stderr)

-        print(traceback.format_exc(), file=sys.stderr)

+        errors.report("Error making Real-ESRGAN models list", exc_info=True)