File size: 6,009 Bytes
e49a8d7 0f528f7 e49a8d7 0f528f7 e49a8d7 0f528f7 e49a8d7 0f528f7 e49a8d7 0f528f7 e49a8d7 0f528f7 e49a8d7 0f528f7 e49a8d7 0f528f7 e49a8d7 cb0572c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 |
import os
import gradio as gr
import torch
from basicsr.archs.rrdbnet_arch import RRDBNet
from basicsr.archs.srvgg_arch import SRVGGNetCompact
from gfpgan.utils import GFPGANer
from huggingface_hub import hf_hub_download
from realesrgan.utils import RealESRGANer
REALESRGAN_REPO_ID = 'leonelhs/realesrgan'
GFPGAN_REPO_ID = 'leonelhs/gfpgan'
os.system("pip freeze")
def showGPU():
if torch.cuda.is_available():
devices = torch.cuda.device_count()
current = torch.cuda.current_device()
return f"Running on GPU:{current} of {devices} total devices"
return "Running on CPU"
def download_model_gfpgan(file):
return hf_hub_download(repo_id=GFPGAN_REPO_ID, filename=file)
def download_model_realesrgan(file):
return hf_hub_download(repo_id=REALESRGAN_REPO_ID, filename=file)
def select_upsampler(version, netscale=4):
model = None
dni_weight = None
version = version + ".pth"
model_path = download_model_realesrgan(version)
if version == 'RealESRGAN_x4plus.pth': # 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)
if version == 'RealESRNet_x4plus.pth': # 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)
if version == 'AI-Forever_x4plus.pth': # 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)
if version == 'RealESRGAN_x4plus_anime_6B.pth': # 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)
if version == 'RealESRGAN_x2plus.pth': # 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 # This is
if version == 'AI-Forever_x2plus.pth': # 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 # This is
if version == 'realesr-animevideov3.pth': # 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')
if version == 'realesr-general-x4v3.pth': # 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')
model_path = [
download_model_realesrgan("realesr-general-x4v3.pth"),
download_model_realesrgan("realesr-general-wdn-x4v3.pth")
]
dni_weight = [0.2, 0.8]
half = True if torch.cuda.is_available() else False
return RealESRGANer(
scale=netscale,
model_path=model_path,
dni_weight=dni_weight,
model=model,
tile=0,
tile_pad=10,
pre_pad=0,
half=half,
gpu_id=0)
def select_face_enhancer(version, scale, upsampler):
if 'v1.2' in version:
model_path = download_model_gfpgan('GFPGANv1.2.pth')
return GFPGANer(
model_path=model_path, upscale=scale, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
elif 'v1.3' in version:
model_path = download_model_gfpgan('GFPGANv1.3.pth')
return GFPGANer(
model_path=model_path, upscale=scale, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
elif 'v1.4' in version:
model_path = download_model_gfpgan('GFPGANv1.4.pth')
return GFPGANer(
model_path=model_path, upscale=scale, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
elif 'RestoreFormer' in version:
model_path = download_model_gfpgan('RestoreFormer.pth')
return GFPGANer(
model_path=model_path, upscale=scale, arch='RestoreFormer', channel_multiplier=2, bg_upsampler=upsampler)
def predict(image, version_upsampler, version_enhancer, scale):
scale = int(scale)
upsampler = select_upsampler(version_upsampler)
if "No additional" not in version_enhancer:
face_enhancer = select_face_enhancer(version_enhancer, scale, upsampler)
_, _, output = face_enhancer.enhance(image, has_aligned=False, only_center_face=False, paste_back=True)
else:
output, _ = upsampler.enhance(image, outscale=scale)
log = f"General enhance version: {version_upsampler}\n " \
f"Face enhance version: {version_enhancer} \n " \
f"Scale:{scale} \n {showGPU()}"
return output, log
title = "Super Face"
description = r"""
<b>Practical Image Restoration Algorithm based on Real-ESRGAN, GFPGAN</b>
"""
article = r"""
<center><span>[email protected] or [email protected]</span></center>
</br>
<center><a href='https://github.com/TencentARC/GFPGAN' target='_blank'>Github Repo ⭐ </a> are welcome</center>
"""
demo = gr.Interface(
predict, [
gr.Image(type="numpy", label="Input"),
gr.Dropdown([
'RealESRGAN_x2plus',
'RealESRGAN_x4plus',
'RealESRNet_x4plus',
'AI-Forever_x2plus',
'AI-Forever_x4plus',
'RealESRGAN_x4plus_anime_6B',
'realesr-animevideov3',
'realesr-general-x4v3'],
type="value", value='RealESRGAN_x4plus', label='General restoration algorithm', info="version"),
gr.Dropdown([
'No additional face process',
'GFPGANv1.2',
'GFPGANv1.3',
'GFPGANv1.4',
'RestoreFormer'],
type="value", value='No additional face process', label='Special face restoration algorithm',
info="version"),
gr.Dropdown(["1", "2", "3", "4"], value="2", label="Rescaling factor")
], [
gr.Image(type="numpy", label="Output", interactive=False),
gr.Textbox(label="log info")
],
title=title,
description=description,
article=article)
demo.queue().launch(share=False, debug=True)
|