Add new models, fix shared opts issues

Add General x4x3, GeneralWDN x4x3, and AnimeVideo models from newer ESRGAN releases.

Fix issues caused by renaming ESRGAN_tille values to GAN_tile without using an IDE...
This commit is contained in:
d8ahazard 2022-09-20 20:11:53 -05:00 committed by AUTOMATIC1111
parent cc287d53a5
commit dd5566814a
2 changed files with 29 additions and 8 deletions

View file

@ -92,10 +92,10 @@ def upscale_without_tiling(model, img):
def esrgan_upscale(model, img): def esrgan_upscale(model, img):
if opts.ESRGAN_tile == 0: if opts.GAN_tile == 0:
return upscale_without_tiling(model, img) return upscale_without_tiling(model, img)
grid = modules.images.split_grid(img, opts.ESRGAN_tile, opts.ESRGAN_tile, opts.ESRGAN_tile_overlap) grid = modules.images.split_grid(img, opts.GAN_tile, opts.GAN_tile, opts.GAN_tile_overlap)
newtiles = [] newtiles = []
scale_factor = 1 scale_factor = 1

View file

@ -2,7 +2,11 @@ import sys
import traceback import traceback
from collections import namedtuple from collections import namedtuple
import numpy as np import numpy as np
import torch
from PIL import Image from PIL import Image
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer
from realesrgan.archs.srvgg_arch import SRVGGNetCompact
import modules.images import modules.images
from modules.shared import cmd_opts, opts from modules.shared import cmd_opts, opts
@ -35,6 +39,24 @@ def setup_realesrgan():
from realesrgan.archs.srvgg_arch import SRVGGNetCompact from realesrgan.archs.srvgg_arch import SRVGGNetCompact
realesrgan_models = [ 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( RealesrganModelInfo(
name="Real-ESRGAN 4x plus", name="Real-ESRGAN 4x plus",
location="https://github.com/xinntao/Real-ESRGA N/releases/download/v0.1.0/RealESRGAN_x4plus.pth", location="https://github.com/xinntao/Real-ESRGA N/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
@ -64,21 +86,20 @@ def setup_realesrgan():
realesrgan_models = [RealesrganModelInfo('None', '', 0, None)] realesrgan_models = [RealesrganModelInfo('None', '', 0, None)]
have_realesrgan = False have_realesrgan = False
def upscale_with_realesrgan(image, RealESRGAN_upscaling, RealESRGAN_model_index): def upscale_with_realesrgan(image, RealESRGAN_upscaling, RealESRGAN_model_index):
if not have_realesrgan or RealESRGANer_constructor is None: if not have_realesrgan:
return image return image
info = realesrgan_models[RealESRGAN_model_index] info = realesrgan_models[RealESRGAN_model_index]
model = info.model() model = info.model()
upsampler = RealESRGANer_constructor( upsampler = RealESRGANer(
scale=info.netscale, scale=info.netscale,
model_path=info.location, model_path=info.location,
model=model, model=model,
half=not cmd_opts.no_half, half=not cmd_opts.no_half,
tile=opts.ESRGAN_tile, tile=opts.GAN_tile,
tile_pad=opts.ESRGAN_tile_overlap, tile_pad=opts.GAN_tile_overlap,
) )
upsampled = upsampler.enhance(np.array(image), outscale=RealESRGAN_upscaling)[0] upsampled = upsampler.enhance(np.array(image), outscale=RealESRGAN_upscaling)[0]