Commit dd556681 authored by d8ahazard's avatar d8ahazard Committed by AUTOMATIC1111

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...
parent cc287d53
...@@ -92,10 +92,10 @@ def upscale_without_tiling(model, img): ...@@ -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
......
...@@ -2,7 +2,11 @@ import sys ...@@ -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,9 +39,27 @@ def setup_realesrgan(): ...@@ -35,9 +39,27 @@ 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-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth", location="https://github.com/xinntao/Real-ESRGA N/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) 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( RealesrganModelInfo(
...@@ -64,21 +86,20 @@ def setup_realesrgan(): ...@@ -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]
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment