diff options
author | d8ahazard <d8ahazard@gmail.com> | 2022-09-21 17:23:38 -0500 |
---|---|---|
committer | AUTOMATIC1111 <16777216c@gmail.com> | 2022-09-22 07:21:54 +0300 |
commit | c214c428bcb28b8b32cbf42dc56eaf5e4cb17f06 (patch) | |
tree | 1ec7618015e01ed8968c319cca0dbc5e916d8c3b /modules/realesrgan_model_loader.py | |
parent | 4f434c8aea3007cbf8038aec0c2a73c452ec9eca (diff) |
Basic settings for realesrgan model selection.
I don't like that you have to restart the app, but it works.
Diffstat (limited to 'modules/realesrgan_model_loader.py')
-rw-r--r-- | modules/realesrgan_model_loader.py | 62 |
1 files changed, 62 insertions, 0 deletions
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) |