Merge branch 'master' into fix-vram

1 files changed, 13 insertions, 6 deletions

diff --git a/modules/gfpgan_model.py b/modules/gfpgan_model.py
index a5fd9632..a9452dce 100644
--- a/modules/gfpgan_model.py
+++ b/modules/gfpgan_model.py
@@ -21,7 +21,7 @@ def gfpgann():
     global loaded_gfpgan_model

     global model_path

     if loaded_gfpgan_model is not None:

-        loaded_gfpgan_model.gfpgan.to(shared.device)

+        loaded_gfpgan_model.gfpgan.to(devices.device_gfpgan)

         return loaded_gfpgan_model

 

     if gfpgan_constructor is None:

@@ -37,25 +37,32 @@ def gfpgann():
         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)

-    model.gfpgan.to(shared.device)

     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):

-    global loaded_gfpgan_model

     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:

-        del model

-        loaded_gfpgan_model = None

-        devices.torch_gc()

+        send_model_to(model, devices.cpu)

 

     return np_image