Spaces:
Runtime error
Runtime error
File size: 1,655 Bytes
71c9afb |
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 |
import os
import torch
import gfpgan
from PIL import Image
from upscaler.RealESRGAN import RealESRGAN
face_enhancer_list = ['NONE', 'GFPGAN', 'REAL-ESRGAN 2x', 'REAL-ESRGAN 4x', 'REAL-ESRGAN 8x']
def load_face_enhancer_model(name='GFPGAN', device="cpu"):
if name == 'GFPGAN':
model_path = "./assets/pretrained_models/GFPGANv1.4.pth"
model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), model_path)
model = gfpgan.GFPGANer(model_path=model_path, upscale=1)
elif name == 'REAL-ESRGAN 2x':
model_path = "./assets/pretrained_models/RealESRGAN_x2.pth"
model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), model_path)
model = RealESRGAN(device, scale=2)
model.load_weights(model_path, download=False)
elif name == 'REAL-ESRGAN 4x':
model_path = "./assets/pretrained_models/RealESRGAN_x4.pth"
model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), model_path)
model = RealESRGAN(device, scale=4)
model.load_weights(model_path, download=False)
elif name == 'REAL-ESRGAN 8x':
model_path = "./assets/pretrained_models/RealESRGAN_x8.pth"
model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), model_path)
model = RealESRGAN(device, scale=8)
model.load_weights(model_path, download=False)
else:
model = None
return model
def gfpgan_enhance(img, model, has_aligned=True):
_, imgs, _ = model.enhance(img, paste_back=True, has_aligned=has_aligned)
return imgs[0]
def realesrgan_enhance(img, model):
img = model.predict(img)
return img |