add resrgan 8x, allow use 1x and up to 8x extra models, move BSRGAN model, add nearest

4 files changed, 33 insertions, 6 deletions

diff --git a/modules/esrgan_model.py b/modules/esrgan_model.py
index a13cf6ac..c61669b4 100644
--- a/modules/esrgan_model.py
+++ b/modules/esrgan_model.py
@@ -50,6 +50,7 @@ def mod2normal(state_dict):
 def resrgan2normal(state_dict, nb=23):

     # this code is copied from https://github.com/victorca25/iNNfer

     if "conv_first.weight" in state_dict and "body.0.rdb1.conv1.weight" in state_dict:

+        re8x = 0

         crt_net = {}

         items = []

         for k, v in state_dict.items():

@@ -75,10 +76,18 @@ def resrgan2normal(state_dict, nb=23):
         crt_net['model.3.bias'] = state_dict['conv_up1.bias']

         crt_net['model.6.weight'] = state_dict['conv_up2.weight']

         crt_net['model.6.bias'] = state_dict['conv_up2.bias']

-        crt_net['model.8.weight'] = state_dict['conv_hr.weight']

-        crt_net['model.8.bias'] = state_dict['conv_hr.bias']

-        crt_net['model.10.weight'] = state_dict['conv_last.weight']

-        crt_net['model.10.bias'] = state_dict['conv_last.bias']

+

+        if 'conv_up3.weight' in state_dict:

+            # modification supporting: https://github.com/ai-forever/Real-ESRGAN/blob/main/RealESRGAN/rrdbnet_arch.py

+            re8x = 3

+            crt_net['model.9.weight'] = state_dict['conv_up3.weight']

+            crt_net['model.9.bias'] = state_dict['conv_up3.bias']

+

+        crt_net[f'model.{8+re8x}.weight'] = state_dict['conv_hr.weight']

+        crt_net[f'model.{8+re8x}.bias'] = state_dict['conv_hr.bias']

+        crt_net[f'model.{10+re8x}.weight'] = state_dict['conv_last.weight']

+        crt_net[f'model.{10+re8x}.bias'] = state_dict['conv_last.bias']

+

         state_dict = crt_net

     return state_dict

 

diff --git a/modules/modelloader.py b/modules/modelloader.py
index b0f2f33d..e4a6f8ac 100644
--- a/modules/modelloader.py
+++ b/modules/modelloader.py
@@ -85,6 +85,9 @@ def cleanup_models():
     src_path = os.path.join(root_path, "ESRGAN")
     dest_path = os.path.join(models_path, "ESRGAN")
     move_files(src_path, dest_path)
+    src_path = os.path.join(models_path, "BSRGAN")
+    dest_path = os.path.join(models_path, "ESRGAN")
+    move_files(src_path, dest_path, ".pth")
     src_path = os.path.join(root_path, "gfpgan")
     dest_path = os.path.join(models_path, "GFPGAN")
     move_files(src_path, dest_path)
diff --git a/modules/ui.py b/modules/ui.py
index 5055ca64..47610f5c 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -1059,7 +1059,7 @@ def create_ui(wrap_gradio_gpu_call):
 

                 with gr.Tabs(elem_id="extras_resize_mode"):

                     with gr.TabItem('Scale by'):

-                        upscaling_resize = gr.Slider(minimum=1.0, maximum=4.0, step=0.05, label="Resize", value=2)

+                        upscaling_resize = gr.Slider(minimum=1.0, maximum=8.0, step=0.05, label="Resize", value=4)

                     with gr.TabItem('Scale to'):

                         with gr.Group():

                             with gr.Row():

diff --git a/modules/upscaler.py b/modules/upscaler.py
index 6ab2fb40..83fde7ca 100644
--- a/modules/upscaler.py
+++ b/modules/upscaler.py
@@ -10,6 +10,7 @@ import modules.shared
 from modules import modelloader, shared
 
 LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS)
+NEAREST = (Image.Resampling.NEAREST if hasattr(Image, 'Resampling') else Image.NEAREST)
 from modules.paths import models_path
 
 
@@ -57,7 +58,7 @@ class Upscaler:
         dest_w = img.width * scale
         dest_h = img.height * scale
         for i in range(3):
-            if img.width >= dest_w and img.height >= dest_h:
+            if img.width > dest_w and img.height > dest_h:
                 break
             img = self.do_upscale(img, selected_model)
         if img.width != dest_w or img.height != dest_h:
@@ -120,3 +121,17 @@ class UpscalerLanczos(Upscaler):
         self.name = "Lanczos"
         self.scalers = [UpscalerData("Lanczos", None, self)]
 
+
+class UpscalerNearest(Upscaler):
+    scalers = []
+
+    def do_upscale(self, img, selected_model=None):
+        return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=NEAREST)
+
+    def load_model(self, _):
+        pass
+
+    def __init__(self, dirname=None):
+        super().__init__(False)
+        self.name = "Nearest"
+        self.scalers = [UpscalerData("Nearest", None, self)]
\ No newline at end of file