reworking #775

3 files changed, 59 insertions, 78 deletions

diff --git a/modules/realesrgan_model.py b/modules/realesrgan_model.py
index fc8623de..c32d6c4c 100644
--- a/modules/realesrgan_model.py
+++ b/modules/realesrgan_model.py
@@ -7,13 +7,60 @@ 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

+

+

+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("Error makeing Real-ESRGAN midels list:", file=sys.stderr)

+        print(traceback.format_exc(), file=sys.stderr)

 

 

 class UpscalerRealESRGAN(modules.images.Upscaler):

@@ -29,16 +76,14 @@ class UpscalerRealESRGAN(modules.images.Upscaler):
 def setup_realesrgan():

     global realesrgan_models

     global have_realesrgan

-    global RealESRGANer_constructor

 

     try:

         from basicsr.archs.rrdbnet_arch import RRDBNet

         from realesrgan import RealESRGANer

         from realesrgan.archs.srvgg_arch import SRVGGNetCompact

 

-        realesrgan_models = realesrgan_model_loader.get_realesrgan_models()

+        realesrgan_models = get_realesrgan_models()

         have_realesrgan = True

-        RealESRGANer_constructor = RealESRGANer

 

         for i, model in enumerate(realesrgan_models):

             if model.name in opts.realesrgan_enabled_models:

diff --git a/modules/realesrgan_model_loader.py b/modules/realesrgan_model_loader.py
deleted file mode 100644
index 3ae3f8b7..00000000
--- a/modules/realesrgan_model_loader.py
+++ /dev/null
@@ -1,62 +0,0 @@
-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 99762f41..0a852838 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -13,7 +13,6 @@ 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

@@ -96,6 +95,11 @@ face_restorers = []
 modules.sd_models.list_models()

 

 

+def realesrgan_models_names():

+    import modules.realesrgan_model

+    return [x.name for x in modules.realesrgan_model.get_realesrgan_models()]

+

+

 class Options:

     class OptionInfo:

         def __init__(self, default=None, label="", component=None, component_args=None, onchange=None):

@@ -143,18 +147,12 @@ 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()}),

+        "realesrgan_enabled_models": OptionInfo(["Real-ESRGAN 4x plus", "Real-ESRGAN 4x plus anime 6B"],"Select which RealESRGAN models to show in the web UI. (Requires restart)", gr.CheckboxGroup, lambda: {"choices": realesrgan_models_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",

-                                          gr.Slider, {"minimum": 1, "maximum": 200, "step": 1}),

-        "ldsr_pre_down":OptionInfo(1, "LDSR Pre-process downssample scale. 1 = no down-sampling, 4 = 1/4 scale.",

-                                          gr.Slider, {"minimum": 1, "maximum": 4, "step": 1}),

-        "ldsr_post_down":OptionInfo(1, "LDSR Post-process down-sample scale. 1 = no down-sampling, 4 = 1/4 scale.",

-                                          gr.Slider, {"minimum": 1, "maximum": 4, "step": 1}),

+        "ldsr_steps": OptionInfo(100, "LDSR processing steps. Lower = faster", gr.Slider, {"minimum": 1, "maximum": 200, "step": 1}),

+        "ldsr_pre_down":OptionInfo(1, "LDSR Pre-process downssample scale. 1 = no down-sampling, 4 = 1/4 scale.", gr.Slider, {"minimum": 1, "maximum": 4, "step": 1}),

+        "ldsr_post_down":OptionInfo(1, "LDSR Post-process down-sample scale. 1 = no down-sampling, 4 = 1/4 scale.", gr.Slider, {"minimum": 1, "maximum": 4, "step": 1}),

         "random_artist_categories": OptionInfo([], "Allowed categories for random artists selection when using the Roll button", gr.CheckboxGroup, {"choices": artist_db.categories()}),

         "upscale_at_full_resolution_padding": OptionInfo(16, "Inpainting at full resolution: padding, in pixels, for the masked region.", gr.Slider, {"minimum": 0, "maximum": 128, "step": 4}),

         "upscaler_for_hires_fix": OptionInfo(None, "Upscaler for highres. fix", gr.Radio, lambda: {"choices": [x.name for x in sd_upscalers]}),