Spaces:
Running
Running
import os | |
import torch | |
from basicsr.utils.download_util import load_file_from_url | |
from basicsr.archs.rrdbnet_arch import RRDBNet | |
from basicsr.archs.srvgg_arch import SRVGGNetCompact | |
from gfpgan.utils import GFPGANer | |
from realesrgan.utils import RealESRGANer | |
from config import * | |
from srcnn import SRCNN | |
def get_upsampler(model_name, device=None): | |
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", | |
] | |
elif model_name == "srcnn": | |
model = SRCNN(device=device) | |
model_path = os.path.join(ROOT_DIR, WEIGHT_DIR, model_name + ".pth") | |
model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu"))) | |
if device: | |
model.to(device) | |
return model | |
else: | |
raise ValueError(f"Wrong model version {model_name}.") | |
model_path = os.path.join(ROOT_DIR, WEIGHT_DIR, model_name + ".pth") | |
if not os.path.exists(model_path): | |
print(f"Downloading weights for model {model_name}") | |
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, WEIGHT_DIR), | |
progress=True, | |
file_name=None, | |
) | |
if model_name != "realesr-general-x4v3": | |
dni_weight = None | |
else: | |
dni_weight = [0.5, 0.5] | |
wdn_model_path = model_path.replace( | |
"realesr-general-x4v3", "realesr-general-wdn-x4v3" | |
) | |
model_path = [model_path, wdn_model_path] | |
half = "cuda" in str(device) | |
return RealESRGANer( | |
scale=netscale, | |
model_path=model_path, | |
dni_weight=dni_weight, | |
model=model, | |
half=half, | |
device=device, | |
) | |
def get_face_enhancer(model_name, upscale=2, bg_upsampler=None, device=None): | |
if model_name == "GFPGANv1.3": | |
arch = "clean" | |
channel_multiplier = 2 | |
file_url = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth" | |
elif model_name == "GFPGANv1.4": | |
arch = "clean" | |
channel_multiplier = 2 | |
file_url = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth" | |
elif model_name == "RestoreFormer": | |
arch = "RestoreFormer" | |
channel_multiplier = 2 | |
file_url = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/RestoreFormer.pth" | |
else: | |
raise ValueError(f"Wrong model version {model_name}.") | |
model_path = os.path.join(ROOT_DIR, WEIGHT_DIR, model_name + ".pth") | |
if not os.path.exists(model_path): | |
print(f"Downloading weights for model {model_name}") | |
model_path = load_file_from_url( | |
url=file_url, | |
model_dir=os.path.join(ROOT_DIR, WEIGHT_DIR), | |
progress=True, | |
file_name=None, | |
) | |
return GFPGANer( | |
model_path=model_path, | |
upscale=upscale, | |
arch=arch, | |
channel_multiplier=channel_multiplier, | |
bg_upsampler=bg_upsampler, | |
device=device, | |
) | |