apextoaster
/
onnx-web
1
0
Fork 0
Code Actions

fix(api): convert Real ESRGAN v3 using same arch as runtime

This commit is contained in:
Sean Sube 2023-02-18 08:41:15 -06:00
parent 2c9d96d2ee
commit 338fc237c7
Signed by: ssube
GPG Key ID: 3EED7B957D362AF1
2 changed files with 28 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
api/onnx_web/chain/upscale_resrgan.py
View File

@ -14,7 +14,7 @@ logger = getLogger(__name__)
last_pipeline_instance = None last_pipeline_instance = None
last_pipeline_params = (None, None) last_pipeline_params = (None, None)
x4_v3_tag = "real-esrgan-x4-v3" TAG_X4_V3 = "real-esrgan-x4-v3"
def load_resrgan( def load_resrgan(
@ -37,7 +37,7 @@ def load_resrgan(
if not path.isfile(model_path): if not path.isfile(model_path):
raise Exception("Real ESRGAN model not found at %s" % model_path) raise Exception("Real ESRGAN model not found at %s" % model_path)
elif params.format == "onnx": if params.format == "onnx":
# use ONNX acceleration, if available # use ONNX acceleration, if available
model = OnnxNet( model = OnnxNet(
server, server,
@ -46,7 +46,7 @@ def load_resrgan(
sess_options=device.sess_options(), sess_options=device.sess_options(),
) )
elif params.format == "pth": elif params.format == "pth":
if x4_v3_tag in model_file: if TAG_X4_V3 in model_file:
# the x4-v3 model needs a different network # the x4-v3 model needs a different network
model = SRVGGNetCompact( model = SRVGGNetCompact(
num_in_ch=3, num_in_ch=3,
@ -69,8 +69,8 @@ def load_resrgan(
raise Exception("unknown platform %s" % params.format) raise Exception("unknown platform %s" % params.format)
dni_weight = None dni_weight = None
if params.upscale_model == x4_v3_tag and params.denoise != 1: if params.upscale_model == TAG_X4_V3 and params.denoise != 1:
wdn_model_path = model_path.replace(x4_v3_tag, "%s-wdn" % (x4_v3_tag)) wdn_model_path = model_path.replace(TAG_X4_V3, "%s-wdn" % TAG_X4_V3)
model_path = [model_path, wdn_model_path] model_path = [model_path, wdn_model_path]
dni_weight = [params.denoise, 1 - params.denoise] dni_weight = [params.denoise, 1 - params.denoise]

31
api/onnx_web/convert/upscale_resrgan.py
View File

@ -3,12 +3,15 @@ from os import path
import torch import torch
from basicsr.archs.rrdbnet_arch import RRDBNet from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan.archs.srvgg_arch import SRVGGNetCompact
from torch.onnx import export from torch.onnx import export
from .utils import ConversionContext, ModelDict from .utils import ConversionContext, ModelDict
logger = getLogger(__name__) logger = getLogger(__name__)
TAG_X4_V3 = "real-esrgan-x4-v3"
@torch.no_grad() @torch.no_grad()
def convert_upscale_resrgan( def convert_upscale_resrgan(
@ -28,14 +31,26 @@ def convert_upscale_resrgan(
return return
logger.info("loading and training model") logger.info("loading and training model")
model = RRDBNet(
num_in_ch=3, if TAG_X4_V3 in name:
num_out_ch=3, # the x4-v3 model needs a different network
num_feat=64, model = SRVGGNetCompact(
num_block=23, num_in_ch=3,
num_grow_ch=32, num_out_ch=3,
scale=scale, num_feat=64,
) num_conv=32,
upscale=scale,
act_type="prelu",
)
else:
model = RRDBNet(
num_in_ch=3,
num_out_ch=3,
num_feat=64,
num_block=23,
num_grow_ch=32,
scale=scale,
)
torch_model = torch.load(source, map_location=ctx.map_location) torch_model = torch.load(source, map_location=ctx.map_location)
if "params_ema" in torch_model: if "params_ema" in torch_model: