WebUI/modules/gfpgan_model.py

118 lines
3.8 KiB
Python
Raw Normal View History

import os
import sys
import traceback
import facexlib
import gfpgan
2022-09-07 09:32:28 +00:00
import modules.face_restoration
2022-09-26 14:29:50 +00:00
from modules import shared, devices, modelloader
from modules.paths import models_path
2022-09-26 14:29:50 +00:00
model_dir = "GFPGAN"
user_path = None
2022-09-26 14:29:50 +00:00
model_path = os.path.join(models_path, model_dir)
model_url = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"
have_gfpgan = False
loaded_gfpgan_model = None
def gfpgann():
global loaded_gfpgan_model
2022-09-26 14:29:50 +00:00
global model_path
2022-09-03 14:28:30 +00:00
if loaded_gfpgan_model is not None:
loaded_gfpgan_model.gfpgan.to(devices.device_gfpgan)
2022-09-03 14:28:30 +00:00
return loaded_gfpgan_model
2022-09-03 14:28:30 +00:00
if gfpgan_constructor is None:
return None
models = modelloader.load_models(model_path, model_url, user_path, ext_filter="GFPGAN")
if len(models) == 1 and "http" in models[0]:
model_file = models[0]
elif len(models) != 0:
2022-09-26 14:29:50 +00:00
latest_file = max(models, key=os.path.getctime)
model_file = latest_file
else:
print("Unable to load gfpgan model!")
return None
if hasattr(facexlib.detection.retinaface, 'device'):
facexlib.detection.retinaface.device = devices.device_gfpgan
model = gfpgan_constructor(model_path=model_file, upscale=1, arch='clean', channel_multiplier=2, bg_upsampler=None, device=devices.device_gfpgan)
loaded_gfpgan_model = model
2022-09-03 14:28:30 +00:00
return model
def send_model_to(model, device):
model.gfpgan.to(device)
model.face_helper.face_det.to(device)
model.face_helper.face_parse.to(device)
def gfpgan_fix_faces(np_image):
model = gfpgann()
2022-09-26 14:29:50 +00:00
if model is None:
return np_image
send_model_to(model, devices.device_gfpgan)
np_image_bgr = np_image[:, :, ::-1]
cropped_faces, restored_faces, gfpgan_output_bgr = model.enhance(np_image_bgr, has_aligned=False, only_center_face=False, paste_back=True)
np_image = gfpgan_output_bgr[:, :, ::-1]
model.face_helper.clean_all()
if shared.opts.face_restoration_unload:
send_model_to(model, devices.cpu)
return np_image
gfpgan_constructor = None
2022-09-26 14:29:50 +00:00
def setup_model(dirname):
global model_path
if not os.path.exists(model_path):
os.makedirs(model_path)
2022-09-26 14:29:50 +00:00
try:
from gfpgan import GFPGANer
from facexlib import detection, parsing
global user_path
global have_gfpgan
global gfpgan_constructor
2022-09-26 14:29:50 +00:00
load_file_from_url_orig = gfpgan.utils.load_file_from_url
facex_load_file_from_url_orig = facexlib.detection.load_file_from_url
facex_load_file_from_url_orig2 = facexlib.parsing.load_file_from_url
def my_load_file_from_url(**kwargs):
return load_file_from_url_orig(**dict(kwargs, model_dir=model_path))
def facex_load_file_from_url(**kwargs):
return facex_load_file_from_url_orig(**dict(kwargs, save_dir=model_path, model_dir=None))
def facex_load_file_from_url2(**kwargs):
return facex_load_file_from_url_orig2(**dict(kwargs, save_dir=model_path, model_dir=None))
gfpgan.utils.load_file_from_url = my_load_file_from_url
facexlib.detection.load_file_from_url = facex_load_file_from_url
facexlib.parsing.load_file_from_url = facex_load_file_from_url2
user_path = dirname
2022-09-26 14:29:50 +00:00
have_gfpgan = True
gfpgan_constructor = GFPGANer
2022-09-07 09:32:28 +00:00
class FaceRestorerGFPGAN(modules.face_restoration.FaceRestoration):
def name(self):
return "GFPGAN"
def restore(self, np_image):
2022-10-03 05:53:52 +00:00
return gfpgan_fix_faces(np_image)
2022-09-07 09:32:28 +00:00
shared.face_restorers.append(FaceRestorerGFPGAN())
except Exception:
print("Error setting up GFPGAN:", file=sys.stderr)
print(traceback.format_exc(), file=sys.stderr)