aboutsummaryrefslogtreecommitdiffstats
path: root/modules/gfpgan_model.py
diff options
context:
space:
mode:
authorbrkirch <brkirch@users.noreply.github.com>2022-10-04 11:42:53 +0000
committerGitHub <noreply@github.com>2022-10-04 11:42:53 +0000
commite9e2a7ec9ac704f133f586eb34176e388c93c87c (patch)
tree969e0e595bd36987ae9de9ae302085ef555bba15 /modules/gfpgan_model.py
parentdc9c5a97742e3a34d37da7108642d8adc0dc5858 (diff)
parentd5bba20a58f43a9f984bb67b4e17f48661f6b818 (diff)
downloadstable-diffusion-webui-gfx803-e9e2a7ec9ac704f133f586eb34176e388c93c87c.tar.gz
stable-diffusion-webui-gfx803-e9e2a7ec9ac704f133f586eb34176e388c93c87c.tar.bz2
stable-diffusion-webui-gfx803-e9e2a7ec9ac704f133f586eb34176e388c93c87c.zip
Merge branch 'master' into cpu-cmdline-opt
Diffstat (limited to 'modules/gfpgan_model.py')
-rw-r--r--modules/gfpgan_model.py16
1 files changed, 13 insertions, 3 deletions
diff --git a/modules/gfpgan_model.py b/modules/gfpgan_model.py
index f1a564b7..a9452dce 100644
--- a/modules/gfpgan_model.py
+++ b/modules/gfpgan_model.py
@@ -36,23 +36,33 @@ def gfpgann():
else:
print("Unable to load gfpgan model!")
return None
- model = gfpgan_constructor(model_path=model_file, upscale=1, arch='clean', channel_multiplier=2, bg_upsampler=None, device=devices.device_gfpgan)
- model.gfpgan.to(devices.device_gfpgan)
+ model = gfpgan_constructor(model_path=model_file, upscale=1, arch='clean', channel_multiplier=2, bg_upsampler=None)
loaded_gfpgan_model = model
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()
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:
- model.gfpgan.to(devices.cpu)
+ send_model_to(model, devices.cpu)
return np_image