Basic settings for realesrgan model selection.

I don't like that you have to restart the app, but it works.

3 files changed, 72 insertions, 37 deletions

diff --git a/modules/realesrgan_model.py b/modules/realesrgan_model.py
index a3917773..fc8623de 100644
--- a/modules/realesrgan_model.py
+++ b/modules/realesrgan_model.py
@@ -7,10 +7,10 @@ from PIL import Image
 from realesrgan import RealESRGANer

 

 import modules.images

+from modules import realesrgan_model_loader

 from modules.shared import cmd_opts, opts

 

 RealesrganModelInfo = namedtuple("RealesrganModelInfo", ["name", "location", "model", "netscale"])

-

 realesrgan_models = []

 have_realesrgan = False

 RealESRGANer_constructor = None

@@ -36,46 +36,13 @@ def setup_realesrgan():
         from realesrgan import RealESRGANer

         from realesrgan.archs.srvgg_arch import SRVGGNetCompact

 

-        realesrgan_models = [

-            RealesrganModelInfo(

-                name="Real-ESRGAN General x4x3",

-                location="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth",

-                netscale=4,

-                model=lambda: SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')

-            ),

-            RealesrganModelInfo(

-                name="Real-ESRGAN General WDN x4x3",

-                location="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth",

-                netscale=4,

-                model=lambda: SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')

-            ),

-            RealesrganModelInfo(

-                name="Real-ESRGAN AnimeVideo",

-                location="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth",

-                netscale=4,

-                model=lambda: SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu')

-            ),

-            RealesrganModelInfo(

-                name="Real-ESRGAN 4x plus",

-                location="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",

-                netscale=4, model=lambda: RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)

-            ),

-            RealesrganModelInfo(

-                name="Real-ESRGAN 4x plus anime 6B",

-                location="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth",

-                netscale=4, model=lambda: RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)

-            ),

-            RealesrganModelInfo(

-                name="Real-ESRGAN 2x plus",

-                location="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth",

-                netscale=2, model=lambda: RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)

-            ),

-        ]

+        realesrgan_models = realesrgan_model_loader.get_realesrgan_models()

         have_realesrgan = True

         RealESRGANer_constructor = RealESRGANer

 

         for i, model in enumerate(realesrgan_models):

-            modules.shared.sd_upscalers.append(UpscalerRealESRGAN(model.netscale, i))

+            if model.name in opts.realesrgan_enabled_models:

+                modules.shared.sd_upscalers.append(UpscalerRealESRGAN(model.netscale, i))

 

     except Exception:

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

@@ -84,6 +51,7 @@ def setup_realesrgan():
         realesrgan_models = [RealesrganModelInfo('None', '', 0, None)]

         have_realesrgan = False

 

+

 def upscale_with_realesrgan(image, RealESRGAN_upscaling, RealESRGAN_model_index):

     if not have_realesrgan:

         return image

diff --git a/modules/realesrgan_model_loader.py b/modules/realesrgan_model_loader.py
new file mode 100644
index 00000000..3ae3f8b7
--- /dev/null
+++ b/modules/realesrgan_model_loader.py
@@ -0,0 +1,62 @@
+from collections import namedtuple
+
+RealesrganModelInfo = namedtuple("RealesrganModelInfo", ["name", "location", "model", "netscale"])
+
+
+def get_realesrgan_model_names():
+    names = set()
+    models = get_realesrgan_models()
+    for model in models:
+        names.add(model.name)
+    return sorted(names)
+
+
+def get_realesrgan_models():
+    try:
+        from basicsr.archs.rrdbnet_arch import RRDBNet
+        from realesrgan import RealESRGANer
+        from realesrgan.archs.srvgg_arch import SRVGGNetCompact
+        models = [
+            RealesrganModelInfo(
+                name="Real-ESRGAN General x4x3",
+                location="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth",
+                netscale=4,
+                model=lambda: SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4,
+                                              act_type='prelu')
+            ),
+            RealesrganModelInfo(
+                name="Real-ESRGAN General WDN x4x3",
+                location="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth",
+                netscale=4,
+                model=lambda: SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4,
+                                              act_type='prelu')
+            ),
+            RealesrganModelInfo(
+                name="Real-ESRGAN AnimeVideo",
+                location="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth",
+                netscale=4,
+                model=lambda: SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4,
+                                              act_type='prelu')
+            ),
+            RealesrganModelInfo(
+                name="Real-ESRGAN 4x plus",
+                location="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
+                netscale=4,
+                model=lambda: RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+            ),
+            RealesrganModelInfo(
+                name="Real-ESRGAN 4x plus anime 6B",
+                location="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth",
+                netscale=4,
+                model=lambda: RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
+            ),
+            RealesrganModelInfo(
+                name="Real-ESRGAN 2x plus",
+                location="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth",
+                netscale=2,
+                model=lambda: RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
+            ),
+        ]
+        return models
+    except Exception as e:
+        print("Exception loading models: %s" % e)
diff --git a/modules/shared.py b/modules/shared.py
index 280d6c4c..99762f41 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -13,6 +13,7 @@ import modules.styles
 import modules.interrogate

 import modules.memmon

 import modules.sd_models

+from modules.realesrgan_model_loader import get_realesrgan_model_names

 

 sd_model_file = os.path.join(script_path, 'model.ckpt')

 default_sd_model_file = sd_model_file

@@ -142,6 +143,10 @@ class Options:
         "save_txt": OptionInfo(False, "Create a text file next to every image with generation parameters."),

         "ESRGAN_tile": OptionInfo(192, "Tile size for ESRGAN upscalers. 0 = no tiling.", gr.Slider, {"minimum": 0, "maximum": 512, "step": 16}),

         "ESRGAN_tile_overlap": OptionInfo(8, "Tile overlap, in pixels for ESRGAN upscalers. Low values = visible seam.", gr.Slider, {"minimum": 0, "maximum": 48, "step": 1}),

+        "realesrgan_enabled_models": OptionInfo(

+            ["Real-ESRGAN 2x plus", "Real-ESRGAN 4x plus", "Real-ESRGAN 4x plus anime 6B"],

+            "Select which RealESRGAN models to show in the web UI. (Requires restart)",

+            gr.CheckboxGroup, {"choices": get_realesrgan_model_names()}),

         "SWIN_tile": OptionInfo(192, "Tile size for all SwinIR.", gr.Slider, {"minimum": 16, "maximum": 512, "step": 16}),

         "SWIN_tile_overlap": OptionInfo(8, "Tile overlap, in pixels for SwinIR. Low values = visible seam.", gr.Slider, {"minimum": 0, "maximum": 48, "step": 1}),

         "ldsr_steps": OptionInfo(100, "LDSR processing steps. Lower = faster",