Upscaler.load_model: don't return None, just use exceptions

This commit is contained in:
Aarni Koskela
2023-05-29 10:38:51 +03:00
parent e3a973a68d
commit bf67a5dcf4
5 changed files with 52 additions and 64 deletions

View File

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