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:
parent
cc287d53a5
commit
dd5566814a
2 changed files with 29 additions and 8 deletions
|
@ -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
|
||||||
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():
|
||||||
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():
|
||||||
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]
|
||||||
|
|
Loading…
Reference in a new issue