Basic settings for realesrgan model selection.
I don't like that you have to restart the app, but it works.
This commit is contained in:
parent
4f434c8aea
commit
c214c428bc
|
@ -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
|
||||
|
|
62
modules/realesrgan_model_loader.py
Normal file
62
modules/realesrgan_model_loader.py
Normal file
|
@ -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)
|
|
@ -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",
|
||||
|
|
Loading…
Reference in New Issue
Block a user