hoelee/stable-diffusion-webui
1
0
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git synced 2025-08-08 05:12:35 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

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

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

40
extensions-builtin/SwinIR/scripts/swinir_model.py
View File

@@ -1,4 +1,4 @@
import os
import sys
import numpy as np
import torch
@@ -7,8 +7,8 @@ from tqdm import tqdm
from modules import modelloader, devices, script_callbacks, shared
from modules.shared import opts, state
from swinir_model_arch import SwinIR as net
from swinir_model_arch_v2 import Swin2SR as net2
from swinir_model_arch import SwinIR
from swinir_model_arch_v2 import Swin2SR
from modules.upscaler import Upscaler, UpscalerData
@@ -36,8 +36,10 @@ class UpscalerSwinIR(Upscaler):
self.scalers = scalers
def do_upscale(self, img, model_file):
model = self.load_model(model_file)
if model is None:
try:
model = self.load_model(model_file)
except Exception as e:
print(f"Failed loading SwinIR model {model_file}: {e}", file=sys.stderr)
return img
model = model.to(device_swinir, dtype=devices.dtype)
img = upscale(img, model)
@@ -56,25 +58,23 @@ class UpscalerSwinIR(Upscaler):
)
else:
filename = path
if filename is None or not os.path.exists(filename):
return None
if filename.endswith(".v2.pth"):
model = net2(
upscale=scale,
in_chans=3,
img_size=64,
window_size=8,
img_range=1.0,
depths=[6, 6, 6, 6, 6, 6],
embed_dim=180,
num_heads=[6, 6, 6, 6, 6, 6],
mlp_ratio=2,
upsampler="nearest+conv",
resi_connection="1conv",
model = Swin2SR(
upscale=scale,
in_chans=3,
img_size=64,
window_size=8,
img_range=1.0,
depths=[6, 6, 6, 6, 6, 6],
embed_dim=180,
num_heads=[6, 6, 6, 6, 6, 6],
mlp_ratio=2,
upsampler="nearest+conv",
resi_connection="1conv",
)
params = None
else:
model = net(
model = SwinIR(
upscale=scale,
in_chans=3,
img_size=64,