add convert scripts
This commit is contained in:
xororz
450
cugan_exporter.py
Normal file
450
cugan_exporter.py
Normal file
@@ -0,0 +1,450 @@
|
||||
import torch
|
||||
from torch import nn as nn
|
||||
from torch.nn import functional as F
|
||||
import os, sys
|
||||
import numpy as np
|
||||
|
||||
root_path = os.path.abspath(".")
|
||||
sys.path.append(root_path)
|
||||
|
||||
|
||||
def q(inp, cache_mode):
|
||||
maxx = inp.max()
|
||||
minn = inp.min()
|
||||
delta = maxx - minn
|
||||
if cache_mode == 2:
|
||||
return (
|
||||
((inp - minn) / delta * 255).round().byte().cpu(),
|
||||
delta,
|
||||
minn,
|
||||
inp.device,
|
||||
)
|
||||
elif cache_mode == 1:
|
||||
return (
|
||||
((inp - minn) / delta * 255).round().byte(),
|
||||
delta,
|
||||
minn,
|
||||
inp.device,
|
||||
)
|
||||
|
||||
|
||||
def dq(inp, if_half, cache_mode, delta, minn, device):
|
||||
if cache_mode == 2:
|
||||
if if_half == True:
|
||||
return inp.to(device).half() / 255 * delta + minn
|
||||
else:
|
||||
return inp.to(device).float() / 255 * delta + minn
|
||||
elif cache_mode == 1:
|
||||
if if_half == True:
|
||||
return inp.half() / 255 * delta + minn
|
||||
else:
|
||||
return inp.float() / 255 * delta + minn
|
||||
|
||||
|
||||
class SEBlock(nn.Module):
|
||||
def __init__(self, in_channels, reduction=8, bias=False):
|
||||
super(SEBlock, self).__init__()
|
||||
self.conv1 = nn.Conv2d(
|
||||
in_channels, in_channels // reduction, 1, 1, 0, bias=bias
|
||||
)
|
||||
self.conv2 = nn.Conv2d(
|
||||
in_channels // reduction, in_channels, 1, 1, 0, bias=bias
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
if "Half" in x.type():
|
||||
x0 = torch.mean(x.float(), dim=(2, 3), keepdim=True).half()
|
||||
else:
|
||||
x0 = torch.mean(x, dim=(2, 3), keepdim=True)
|
||||
x0 = self.conv1(x0)
|
||||
x0 = F.relu(x0, inplace=True)
|
||||
x0 = self.conv2(x0)
|
||||
x0 = torch.sigmoid(x0)
|
||||
x = torch.mul(x, x0)
|
||||
return x
|
||||
|
||||
def forward_mean(self, x, x0):
|
||||
x0 = self.conv1(x0)
|
||||
x0 = F.relu(x0, inplace=True)
|
||||
x0 = self.conv2(x0)
|
||||
x0 = torch.sigmoid(x0)
|
||||
x = torch.mul(x, x0)
|
||||
return x
|
||||
|
||||
|
||||
class UNetConv(nn.Module):
|
||||
def __init__(self, in_channels, mid_channels, out_channels, se):
|
||||
super(UNetConv, self).__init__()
|
||||
self.conv = nn.Sequential(
|
||||
nn.Conv2d(in_channels, mid_channels, 3, 1, 0),
|
||||
nn.LeakyReLU(0.1, inplace=True),
|
||||
nn.Conv2d(mid_channels, out_channels, 3, 1, 0),
|
||||
nn.LeakyReLU(0.1, inplace=True),
|
||||
)
|
||||
if se:
|
||||
self.seblock = SEBlock(out_channels, reduction=8, bias=True)
|
||||
else:
|
||||
self.seblock = None
|
||||
|
||||
def forward(self, x):
|
||||
z = self.conv(x)
|
||||
if self.seblock is not None:
|
||||
z = self.seblock(z)
|
||||
return z
|
||||
|
||||
|
||||
class UNet1(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, deconv):
|
||||
super(UNet1, self).__init__()
|
||||
self.conv1 = UNetConv(in_channels, 32, 64, se=False)
|
||||
self.conv1_down = nn.Conv2d(64, 64, 2, 2, 0)
|
||||
self.conv2 = UNetConv(64, 128, 64, se=True)
|
||||
self.conv2_up = nn.ConvTranspose2d(64, 64, 2, 2, 0)
|
||||
self.conv3 = nn.Conv2d(64, 64, 3, 1, 0)
|
||||
|
||||
if deconv:
|
||||
self.conv_bottom = nn.ConvTranspose2d(64, out_channels, 4, 2, 3)
|
||||
else:
|
||||
self.conv_bottom = nn.Conv2d(64, out_channels, 3, 1, 0)
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
|
||||
nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
|
||||
elif isinstance(m, nn.Linear):
|
||||
nn.init.normal_(m.weight, 0, 0.01)
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
def forward(self, x):
|
||||
x1 = self.conv1(x)
|
||||
x2 = self.conv1_down(x1)
|
||||
x1 = F.pad(x1, (-4, -4, -4, -4))
|
||||
x2 = F.leaky_relu(x2, 0.1, inplace=True)
|
||||
x2 = self.conv2(x2)
|
||||
x2 = self.conv2_up(x2)
|
||||
x2 = F.leaky_relu(x2, 0.1, inplace=True)
|
||||
x3 = self.conv3(x1 + x2)
|
||||
x3 = F.leaky_relu(x3, 0.1, inplace=True)
|
||||
z = self.conv_bottom(x3)
|
||||
return z
|
||||
|
||||
def forward_a(self, x):
|
||||
x1 = self.conv1(x)
|
||||
x2 = self.conv1_down(x1)
|
||||
x1 = F.pad(x1, (-4, -4, -4, -4))
|
||||
x2 = F.leaky_relu(x2, 0.1, inplace=True)
|
||||
x2 = self.conv2.conv(x2)
|
||||
return x1, x2
|
||||
|
||||
def forward_b(self, x1, x2):
|
||||
x2 = self.conv2_up(x2)
|
||||
x2 = F.leaky_relu(x2, 0.1, inplace=True)
|
||||
x3 = self.conv3(x1 + x2)
|
||||
x3 = F.leaky_relu(x3, 0.1, inplace=True)
|
||||
z = self.conv_bottom(x3)
|
||||
return z
|
||||
|
||||
|
||||
class UNet1x3(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, deconv):
|
||||
super(UNet1x3, self).__init__()
|
||||
self.conv1 = UNetConv(in_channels, 32, 64, se=False)
|
||||
self.conv1_down = nn.Conv2d(64, 64, 2, 2, 0)
|
||||
self.conv2 = UNetConv(64, 128, 64, se=True)
|
||||
self.conv2_up = nn.ConvTranspose2d(64, 64, 2, 2, 0)
|
||||
self.conv3 = nn.Conv2d(64, 64, 3, 1, 0)
|
||||
|
||||
if deconv:
|
||||
self.conv_bottom = nn.ConvTranspose2d(64, out_channels, 5, 3, 2)
|
||||
else:
|
||||
self.conv_bottom = nn.Conv2d(64, out_channels, 3, 1, 0)
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
|
||||
nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
|
||||
elif isinstance(m, nn.Linear):
|
||||
nn.init.normal_(m.weight, 0, 0.01)
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
def forward(self, x):
|
||||
x1 = self.conv1(x)
|
||||
x2 = self.conv1_down(x1)
|
||||
x1 = F.pad(x1, (-4, -4, -4, -4))
|
||||
x2 = F.leaky_relu(x2, 0.1, inplace=True)
|
||||
x2 = self.conv2(x2)
|
||||
x2 = self.conv2_up(x2)
|
||||
x2 = F.leaky_relu(x2, 0.1, inplace=True)
|
||||
x3 = self.conv3(x1 + x2)
|
||||
x3 = F.leaky_relu(x3, 0.1, inplace=True)
|
||||
z = self.conv_bottom(x3)
|
||||
return z
|
||||
|
||||
def forward_a(self, x):
|
||||
x1 = self.conv1(x)
|
||||
x2 = self.conv1_down(x1)
|
||||
x1 = F.pad(x1, (-4, -4, -4, -4))
|
||||
x2 = F.leaky_relu(x2, 0.1, inplace=True)
|
||||
x2 = self.conv2.conv(x2)
|
||||
return x1, x2
|
||||
|
||||
def forward_b(self, x1, x2):
|
||||
x2 = self.conv2_up(x2)
|
||||
x2 = F.leaky_relu(x2, 0.1, inplace=True)
|
||||
x3 = self.conv3(x1 + x2)
|
||||
x3 = F.leaky_relu(x3, 0.1, inplace=True)
|
||||
z = self.conv_bottom(x3)
|
||||
return z
|
||||
|
||||
|
||||
class UNet2(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, deconv):
|
||||
super(UNet2, self).__init__()
|
||||
|
||||
self.conv1 = UNetConv(in_channels, 32, 64, se=False)
|
||||
self.conv1_down = nn.Conv2d(64, 64, 2, 2, 0)
|
||||
self.conv2 = UNetConv(64, 64, 128, se=True)
|
||||
self.conv2_down = nn.Conv2d(128, 128, 2, 2, 0)
|
||||
self.conv3 = UNetConv(128, 256, 128, se=True)
|
||||
self.conv3_up = nn.ConvTranspose2d(128, 128, 2, 2, 0)
|
||||
self.conv4 = UNetConv(128, 64, 64, se=True)
|
||||
self.conv4_up = nn.ConvTranspose2d(64, 64, 2, 2, 0)
|
||||
self.conv5 = nn.Conv2d(64, 64, 3, 1, 0)
|
||||
|
||||
if deconv:
|
||||
self.conv_bottom = nn.ConvTranspose2d(64, out_channels, 4, 2, 3)
|
||||
else:
|
||||
self.conv_bottom = nn.Conv2d(64, out_channels, 3, 1, 0)
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
|
||||
nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
|
||||
elif isinstance(m, nn.Linear):
|
||||
nn.init.normal_(m.weight, 0, 0.01)
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
def forward(self, x, alpha=1):
|
||||
x1 = self.conv1(x)
|
||||
x2 = self.conv1_down(x1)
|
||||
x1 = F.pad(x1, (-16, -16, -16, -16))
|
||||
x2 = F.leaky_relu(x2, 0.1, inplace=True)
|
||||
x2 = self.conv2(x2)
|
||||
x3 = self.conv2_down(x2)
|
||||
x2 = F.pad(x2, (-4, -4, -4, -4))
|
||||
x3 = F.leaky_relu(x3, 0.1, inplace=True)
|
||||
x3 = self.conv3(x3)
|
||||
x3 = self.conv3_up(x3)
|
||||
x3 = F.leaky_relu(x3, 0.1, inplace=True)
|
||||
x4 = self.conv4(x2 + x3)
|
||||
x4 *= alpha
|
||||
x4 = self.conv4_up(x4)
|
||||
x4 = F.leaky_relu(x4, 0.1, inplace=True)
|
||||
x5 = self.conv5(x1 + x4)
|
||||
x5 = F.leaky_relu(x5, 0.1, inplace=True)
|
||||
z = self.conv_bottom(x5)
|
||||
return z
|
||||
|
||||
def forward_a(self, x):
|
||||
x1 = self.conv1(x)
|
||||
x2 = self.conv1_down(x1)
|
||||
x1 = F.pad(x1, (-16, -16, -16, -16))
|
||||
x2 = F.leaky_relu(x2, 0.1, inplace=True)
|
||||
x2 = self.conv2.conv(x2)
|
||||
return x1, x2
|
||||
|
||||
def forward_b(self, x2):
|
||||
x3 = self.conv2_down(x2)
|
||||
x2 = F.pad(x2, (-4, -4, -4, -4))
|
||||
x3 = F.leaky_relu(x3, 0.1, inplace=True)
|
||||
x3 = self.conv3.conv(x3)
|
||||
return x2, x3
|
||||
|
||||
def forward_c(self, x2, x3):
|
||||
x3 = self.conv3_up(x3)
|
||||
x3 = F.leaky_relu(x3, 0.1, inplace=True)
|
||||
x4 = self.conv4.conv(x2 + x3)
|
||||
return x4
|
||||
|
||||
def forward_d(self, x1, x4):
|
||||
x4 = self.conv4_up(x4)
|
||||
x4 = F.leaky_relu(x4, 0.1, inplace=True)
|
||||
x5 = self.conv5(x1 + x4)
|
||||
x5 = F.leaky_relu(x5, 0.1, inplace=True)
|
||||
|
||||
z = self.conv_bottom(x5)
|
||||
return z
|
||||
|
||||
|
||||
class UpCunet2x(nn.Module):
|
||||
def __init__(self, in_channels=3, out_channels=3):
|
||||
super(UpCunet2x, self).__init__()
|
||||
self.unet1 = UNet1(in_channels, out_channels, deconv=True)
|
||||
self.unet2 = UNet2(in_channels, out_channels, deconv=False)
|
||||
|
||||
def forward(self, x):
|
||||
x = F.pad(x, (18, 18, 18, 18), "reflect")
|
||||
x = self.unet1(x)
|
||||
x0 = self.unet2(x, 1)
|
||||
x = F.pad(x, (-20, -20, -20, -20))
|
||||
x = torch.add(x0, x)
|
||||
return x
|
||||
|
||||
|
||||
class UpCunet3x(nn.Module):
|
||||
def __init__(self, in_channels=3, out_channels=3):
|
||||
super(UpCunet3x, self).__init__()
|
||||
self.unet1 = UNet1x3(in_channels, out_channels, deconv=True)
|
||||
self.unet2 = UNet2(in_channels, out_channels, deconv=False)
|
||||
|
||||
def forward(self, x):
|
||||
x = F.pad(x, (14, 14, 14, 14), "reflect")
|
||||
x = self.unet1(x)
|
||||
x0 = self.unet2(x, 1)
|
||||
x = F.pad(x, (-20, -20, -20, -20))
|
||||
x = torch.add(x0, x)
|
||||
return x
|
||||
|
||||
|
||||
class UpCunet4x(nn.Module):
|
||||
def __init__(self, in_channels=3, out_channels=3):
|
||||
super(UpCunet4x, self).__init__()
|
||||
self.unet1 = UNet1(in_channels, 64, deconv=True)
|
||||
self.unet2 = UNet2(64, 64, deconv=False)
|
||||
self.ps = nn.PixelShuffle(2)
|
||||
self.conv_final = nn.Conv2d(64, 12, 3, 1, padding=0, bias=True)
|
||||
|
||||
def forward(self, x):
|
||||
x00 = x
|
||||
x = F.pad(x, (19, 19, 19, 19), "reflect")
|
||||
x = self.unet1.forward(x)
|
||||
x0 = self.unet2.forward(x, 1)
|
||||
x1 = F.pad(x, (-20, -20, -20, -20))
|
||||
x = torch.add(x0, x1)
|
||||
x = self.conv_final(x)
|
||||
x = F.pad(x, (-1, -1, -1, -1))
|
||||
x = self.ps(x)
|
||||
x += F.interpolate(x00, scale_factor=4, mode="nearest")
|
||||
return x
|
||||
|
||||
|
||||
class RealWaifuUpScaler(object):
|
||||
def __init__(self, scale, weight_path, half, device):
|
||||
weight = torch.load(weight_path, map_location="cpu")
|
||||
self.pro = "pro" in weight
|
||||
if self.pro:
|
||||
del weight["pro"]
|
||||
self.model = eval("UpCunet%sx" % scale)()
|
||||
|
||||
if half == True:
|
||||
self.model = self.model.half().to(device)
|
||||
else:
|
||||
self.model = self.model.to(device)
|
||||
self.model.load_state_dict(weight, strict=False)
|
||||
self.model.eval()
|
||||
self.half = half
|
||||
self.device = device
|
||||
|
||||
def np2tensor(self, np_frame):
|
||||
if self.pro:
|
||||
if self.half == False:
|
||||
return (
|
||||
torch.from_numpy(np.transpose(np_frame, (2, 0, 1)))
|
||||
.unsqueeze(0)
|
||||
.to(self.device)
|
||||
.float()
|
||||
/ (255 / 0.7)
|
||||
+ 0.15
|
||||
)
|
||||
else:
|
||||
return (
|
||||
torch.from_numpy(np.transpose(np_frame, (2, 0, 1)))
|
||||
.unsqueeze(0)
|
||||
.to(self.device)
|
||||
.half()
|
||||
/ (255 / 0.7)
|
||||
+ 0.15
|
||||
)
|
||||
else:
|
||||
if self.half == False:
|
||||
return (
|
||||
torch.from_numpy(np.transpose(np_frame, (2, 0, 1)))
|
||||
.unsqueeze(0)
|
||||
.to(self.device)
|
||||
.float()
|
||||
/ 255
|
||||
)
|
||||
else:
|
||||
return (
|
||||
torch.from_numpy(np.transpose(np_frame, (2, 0, 1)))
|
||||
.unsqueeze(0)
|
||||
.to(self.device)
|
||||
.half()
|
||||
/ 255
|
||||
)
|
||||
|
||||
def tensor2np(self, tensor):
|
||||
return np.transpose(tensor.squeeze().cpu().numpy(), (1, 2, 0))
|
||||
|
||||
def __call__(self, frame, tile_mode, cache_mode, alpha):
|
||||
with torch.no_grad():
|
||||
tensor = self.np2tensor(frame)
|
||||
if cache_mode == 3:
|
||||
result = self.tensor2np(
|
||||
self.model.forward_gap_sync(tensor, tile_mode, alpha, self.pro)
|
||||
)
|
||||
elif cache_mode == 2:
|
||||
result = self.tensor2np(
|
||||
self.model.forward_fast_rough(tensor, tile_mode, alpha, self.pro)
|
||||
)
|
||||
else:
|
||||
result = self.tensor2np(self.model(tensor))
|
||||
return result
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
for weight_path, scale, name in [
|
||||
("./weights_v3/up2x-latest-conservative.pth", 2, "2x-conservative"),
|
||||
("./weights_v3/up2x-latest-denoise1x.pth", 2, "2x-denoise1x"),
|
||||
("./weights_v3/up2x-latest-denoise2x.pth", 2, "2x-denoise2x"),
|
||||
("./weights_v3/up2x-latest-denoise3x.pth", 2, "2x-denoise3x"),
|
||||
("./weights_v3/up2x-latest-no-denoise.pth", 2, "2x-no-denoise"),
|
||||
("./weights_v3/up3x-latest-conservative.pth", 3, "3x-conservative"),
|
||||
("./weights_v3/up3x-latest-denoise3x.pth", 3, "3x-denoise3x"),
|
||||
# ("./weights_v3/up3x-latest-no-denoise.pth", 3, "3x-no-denoise"), # error
|
||||
("./weights_v3/up4x-latest-conservative.pth", 4, "4x-conservative"),
|
||||
("./weights_v3/up4x-latest-denoise3x.pth", 4, "4x-denoise3x"),
|
||||
("./weights_v3/up4x-latest-no-denoise.pth", 4, "4x-no-denoise"),
|
||||
]:
|
||||
upscaler = RealWaifuUpScaler(scale, weight_path, half=False, device="cpu")
|
||||
# for tile_size in (64, 128):
|
||||
# for tile_size in (32, 48):
|
||||
# for tile_size in (32, 48, 64, 128):
|
||||
# for tile_size in (96,):
|
||||
# for tile_size in (192, 256):
|
||||
for tile_size in (384, 512):
|
||||
torch.onnx.export(
|
||||
upscaler.model,
|
||||
torch.randn(1, 3, tile_size, tile_size),
|
||||
f"{name}.onnx",
|
||||
export_params=True,
|
||||
opset_version=11,
|
||||
do_constant_folding=True,
|
||||
input_names=["input"],
|
||||
output_names=["output"],
|
||||
dynamic_axes={
|
||||
"input": {2: "height", 3: "width"},
|
||||
"output": {2: "height", 3: "width"},
|
||||
},
|
||||
)
|
||||
os.system(f"onnxsim {name}.onnx {name}.onnx")
|
||||
os.system(
|
||||
f"onnx2tf -i {name}.onnx -osd -ois input:1,3,{tile_size},{tile_size} -o {name}-{tile_size}-tf"
|
||||
)
|
||||
# export tfjs_converter = xxx
|
||||
if not os.path.exists("tfjs"):
|
||||
os.mkdir("tfjs")
|
||||
os.system(
|
||||
# f"{os.environ['tfjs_converter']} --quantize_float16 --input_format tf_saved_model --output_format tfjs_graph_model {name}-{tile_size}-tf tfjs/{name}-{tile_size}"
|
||||
f"{os.environ['tfjs_converter']} --control_flow_v2=True --quantize_float16 --input_format tf_saved_model --output_format tfjs_graph_model {name}-{tile_size}-tf tfjs/{name}-{tile_size}"
|
||||
)
|
||||
154
esrgan-exporter.py
Normal file
154
esrgan-exporter.py
Normal file
@@ -0,0 +1,154 @@
|
||||
import os
|
||||
from basicsr.archs.rrdbnet_arch import RRDBNet
|
||||
from basicsr.utils.download_util import load_file_from_url
|
||||
|
||||
from realesrgan import RealESRGANer
|
||||
from realesrgan.archs.srvgg_arch import SRVGGNetCompact
|
||||
import torch
|
||||
|
||||
model2name = {
|
||||
"realesr-animevideov3": "anime_fast",
|
||||
"RealESRGAN_x4plus_anime_6B": "anime_plus",
|
||||
"realesr-general-x4v3": "general_fast",
|
||||
"RealESRGAN_x4plus": "general_plus",
|
||||
}
|
||||
|
||||
|
||||
def export(model_name, tile_size):
|
||||
if model_name == "RealESRGAN_x4plus": # x4 RRDBNet model
|
||||
model = RRDBNet(
|
||||
num_in_ch=3,
|
||||
num_out_ch=3,
|
||||
num_feat=64,
|
||||
num_block=23,
|
||||
num_grow_ch=32,
|
||||
scale=4,
|
||||
)
|
||||
netscale = 4
|
||||
file_url = [
|
||||
"https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"
|
||||
]
|
||||
elif model_name == "RealESRNet_x4plus": # x4 RRDBNet model
|
||||
model = RRDBNet(
|
||||
num_in_ch=3,
|
||||
num_out_ch=3,
|
||||
num_feat=64,
|
||||
num_block=23,
|
||||
num_grow_ch=32,
|
||||
scale=4,
|
||||
)
|
||||
netscale = 4
|
||||
file_url = [
|
||||
"https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth"
|
||||
]
|
||||
elif model_name == "RealESRGAN_x4plus_anime_6B": # x4 RRDBNet model with 6 blocks
|
||||
model = RRDBNet(
|
||||
num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4
|
||||
)
|
||||
netscale = 4
|
||||
file_url = [
|
||||
"https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth"
|
||||
]
|
||||
elif model_name == "RealESRGAN_x2plus": # x2 RRDBNet model
|
||||
model = RRDBNet(
|
||||
num_in_ch=3,
|
||||
num_out_ch=3,
|
||||
num_feat=64,
|
||||
num_block=23,
|
||||
num_grow_ch=32,
|
||||
scale=2,
|
||||
)
|
||||
netscale = 2
|
||||
file_url = [
|
||||
"https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth"
|
||||
]
|
||||
elif model_name == "realesr-animevideov3": # x4 VGG-style model (XS size)
|
||||
model = SRVGGNetCompact(
|
||||
num_in_ch=3,
|
||||
num_out_ch=3,
|
||||
num_feat=64,
|
||||
num_conv=16,
|
||||
upscale=4,
|
||||
act_type="prelu",
|
||||
)
|
||||
netscale = 4
|
||||
file_url = [
|
||||
"https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth"
|
||||
]
|
||||
elif model_name == "realesr-general-x4v3": # x4 VGG-style model (S size)
|
||||
model = SRVGGNetCompact(
|
||||
num_in_ch=3,
|
||||
num_out_ch=3,
|
||||
num_feat=64,
|
||||
num_conv=32,
|
||||
upscale=4,
|
||||
act_type="prelu",
|
||||
)
|
||||
netscale = 4
|
||||
file_url = [
|
||||
"https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth",
|
||||
"https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth",
|
||||
]
|
||||
|
||||
model_path = os.path.join("weights", model_name + ".pth")
|
||||
if not os.path.isfile(model_path):
|
||||
ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
for url in file_url:
|
||||
# model_path will be updated
|
||||
model_path = load_file_from_url(
|
||||
url=url,
|
||||
model_dir=os.path.join(ROOT_DIR, "weights"),
|
||||
progress=True,
|
||||
file_name=None,
|
||||
)
|
||||
|
||||
# use dni to control the denoise strength
|
||||
dni_weight = None
|
||||
denoise_strength = 0.5
|
||||
if model_name == "realesr-general-x4v3" and denoise_strength != 1:
|
||||
wdn_model_path = model_path.replace(
|
||||
"realesr-general-x4v3", "realesr-general-wdn-x4v3"
|
||||
)
|
||||
model_path = [model_path, wdn_model_path]
|
||||
dni_weight = [denoise_strength, 1 - denoise_strength]
|
||||
|
||||
# restorer
|
||||
upsampler = RealESRGANer(
|
||||
scale=netscale,
|
||||
model_path=model_path,
|
||||
dni_weight=dni_weight,
|
||||
model=model,
|
||||
tile=0,
|
||||
tile_pad=0,
|
||||
pre_pad=0,
|
||||
half=False,
|
||||
gpu_id=None,
|
||||
)
|
||||
|
||||
name = model2name[model_name]
|
||||
torch.onnx.export(
|
||||
upsampler.model,
|
||||
torch.randn(1, 3, tile_size, tile_size),
|
||||
f"{name}.onnx",
|
||||
export_params=True,
|
||||
opset_version=11,
|
||||
do_constant_folding=True,
|
||||
input_names=["input"],
|
||||
output_names=["output"],
|
||||
)
|
||||
os.system(f"onnxsim {name}.onnx {name}.onnx")
|
||||
os.system(
|
||||
f"onnx2tf -i {name}.onnx -osd -ois input:1,3,{tile_size},{tile_size} -o {name}-{tile_size}-tf"
|
||||
)
|
||||
# export tfjs_converter = xxx
|
||||
if not os.path.exists("tfjs"):
|
||||
os.mkdir("tfjs")
|
||||
os.system(
|
||||
# f"{os.environ['tfjs_converter']} --quantize_float16 --input_format tf_saved_model --output_format tfjs_graph_model {name}-{tile_size}-tf tfjs/{name}-{tile_size}"
|
||||
f"{os.environ['tfjs_converter']} --control_flow_v2=True --quantize_float16 --input_format tf_saved_model --output_format tfjs_graph_model {name}-{tile_size}-tf tfjs/{name}-{tile_size}"
|
||||
)
|
||||
|
||||
|
||||
for model_name in model2name.keys():
|
||||
for tile_size in [32, 48, 64, 96, 128, 192, 256]:
|
||||
export(model_name, tile_size)
|
||||
Reference in New Issue
Block a user