Merge branch 'dev' into patch-1

1 files changed, 26 insertions, 22 deletions

diff --git a/extensions-builtin/ScuNET/scripts/scunet_model.py b/extensions-builtin/ScuNET/scripts/scunet_model.py
index c7fd5739..ffef26b2 100644
--- a/extensions-builtin/ScuNET/scripts/scunet_model.py
+++ b/extensions-builtin/ScuNET/scripts/scunet_model.py
@@ -1,19 +1,16 @@
-import os.path
 import sys
-import traceback
 
 import PIL.Image
 import numpy as np
 import torch
 from tqdm import tqdm
 
-from basicsr.utils.download_util import load_file_from_url
-
 import modules.upscaler
-from modules import devices, modelloader
-from scunet_model_arch import SCUNet as net
+from modules import devices, modelloader, script_callbacks, errors
+from scunet_model_arch import SCUNet
+
+from modules.modelloader import load_file_from_url
 from modules.shared import opts
-from modules import images
 
 
 class UpscalerScuNET(modules.upscaler.Upscaler):
@@ -29,7 +26,7 @@ class UpscalerScuNET(modules.upscaler.Upscaler):
         scalers = []
         add_model2 = True
         for file in model_paths:
-            if "http" in file:
+            if file.startswith("http"):
                 name = self.model_name
             else:
                 name = modelloader.friendly_name(file)
@@ -39,8 +36,7 @@ class UpscalerScuNET(modules.upscaler.Upscaler):
                 scaler_data = modules.upscaler.UpscalerData(name, file, self, 4)
                 scalers.append(scaler_data)
             except Exception:
-                print(f"Error loading ScuNET model: {file}", file=sys.stderr)
-                print(traceback.format_exc(), file=sys.stderr)
+                errors.report(f"Error loading ScuNET model: {file}", exc_info=True)
         if add_model2:
             scaler_data2 = modules.upscaler.UpscalerData(self.model_name2, self.model_url2, self)
             scalers.append(scaler_data2)
@@ -91,9 +87,10 @@ class UpscalerScuNET(modules.upscaler.Upscaler):
 
         torch.cuda.empty_cache()
 
-        model = self.load_model(selected_file)
-        if model is None:
-            print(f"ScuNET: Unable to load model from {selected_file}", file=sys.stderr)
+        try:
+            model = self.load_model(selected_file)
+        except Exception as e:
+            print(f"ScuNET: Unable to load model from {selected_file}: {e}", file=sys.stderr)
             return img
 
         device = devices.get_device_for('scunet')
@@ -121,20 +118,27 @@ class UpscalerScuNET(modules.upscaler.Upscaler):
 
     def load_model(self, path: str):
         device = devices.get_device_for('scunet')
-        if "http" in path:
-            filename = load_file_from_url(url=self.model_url, model_dir=self.model_path, file_name="%s.pth" % self.name,
-                                          progress=True)
+        if path.startswith("http"):
+            # TODO: this doesn't use `path` at all?
+            filename = load_file_from_url(self.model_url, model_dir=self.model_download_path, file_name=f"{self.name}.pth")
         else:
             filename = path
-        if not os.path.exists(os.path.join(self.model_path, filename)) or filename is None:
-            print(f"ScuNET: Unable to load model from {filename}", file=sys.stderr)
-            return None
-
-        model = net(in_nc=3, config=[4, 4, 4, 4, 4, 4, 4], dim=64)
+        model = SCUNet(in_nc=3, config=[4, 4, 4, 4, 4, 4, 4], dim=64)
         model.load_state_dict(torch.load(filename), strict=True)
         model.eval()
-        for k, v in model.named_parameters():
+        for _, v in model.named_parameters():
             v.requires_grad = False
         model = model.to(device)
 
         return model
+
+
+def on_ui_settings():
+    import gradio as gr
+    from modules import shared
+
+    shared.opts.add_option("SCUNET_tile", shared.OptionInfo(256, "Tile size for SCUNET upscalers.", gr.Slider, {"minimum": 0, "maximum": 512, "step": 16}, section=('upscaling', "Upscaling")).info("0 = no tiling"))
+    shared.opts.add_option("SCUNET_tile_overlap", shared.OptionInfo(8, "Tile overlap for SCUNET upscalers.", gr.Slider, {"minimum": 0, "maximum": 64, "step": 1}, section=('upscaling', "Upscaling")).info("Low values = visible seam"))
+
+
+script_callbacks.on_ui_settings(on_ui_settings)