diff --git a/modules/realesrgan_model.py b/modules/realesrgan_model.py index fc8623ded..c32d6c4c2 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 3ae3f8b77..000000000 --- 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 99762f419..0a8528388 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]}),