From 843b2b64fcd41be4a9e934ba83a3a499c7aff5c0 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Mon, 12 Sep 2022 18:40:06 +0300
Subject:  Instance of CUDA out of memory on a low-res batch, even with
 --opt-split-attention-v1 (found cause) #255

---
 modules/gfpgan_model.py | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

(limited to 'modules/gfpgan_model.py')

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
 
 
-- 
cgit v1.2.1