From 3eaae8ae655346c8c510b677bd48021b86842561 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Thu, 22 Sep 2022 07:59:27 +0300
Subject: reworking #775

---
 modules/realesrgan_model.py | 55 ++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 50 insertions(+), 5 deletions(-)

(limited to 'modules/realesrgan_model.py')

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:
-- 
cgit v1.2.1