Instance of CUDA out of memory on a low-res batch, even with --opt-split-attention-v1 (found cause) #255

1 files changed, 10 insertions, 5 deletions

diff --git a/modules/gfpgan_model.py b/modules/gfpgan_model.py
index f697326c..0af97123 100644
--- a/modules/gfpgan_model.py
+++ b/modules/gfpgan_model.py
@@ -2,7 +2,7 @@ import os
 import sys

 import traceback

 

-from modules import shared

+from modules import shared, devices

 from modules.shared import cmd_opts

 from modules.paths import script_path

 import modules.face_restoration

@@ -28,24 +28,29 @@ def gfpgan():
     global loaded_gfpgan_model

 

     if loaded_gfpgan_model is not None:

+        loaded_gfpgan_model.gfpgan.to(shared.device)

         return loaded_gfpgan_model

 

     if gfpgan_constructor is None:

         return None

 

     model = gfpgan_constructor(model_path=gfpgan_model_path(), upscale=1, arch='clean', channel_multiplier=2, bg_upsampler=None)

-

-    if not cmd_opts.unload_gfpgan:

-        loaded_gfpgan_model = model

+    model.gfpgan.to(shared.device)

+    loaded_gfpgan_model = model

 

     return model

 

 

 def gfpgan_fix_faces(np_image):

+    model = gfpgan()

+

     np_image_bgr = np_image[:, :, ::-1]

-    cropped_faces, restored_faces, gfpgan_output_bgr = gfpgan().enhance(np_image_bgr, has_aligned=False, only_center_face=False, paste_back=True)

+    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]

 

+    if shared.opts.face_restoration_unload:

+        model.gfpgan.to(devices.cpu)

+

     return np_image