option to unload GFPGAN after using

2 files changed, 13 insertions, 3 deletions

diff --git a/modules/gfpgan_model.py b/modules/gfpgan_model.py
index 3f42c163..334a1b7f 100644
--- a/modules/gfpgan_model.py
+++ b/modules/gfpgan_model.py
@@ -4,6 +4,7 @@ import traceback
 

 from modules.paths import script_path

 from modules.shared import cmd_opts

+import modules.shared

 

 

 def gfpgan_model_path():

@@ -23,10 +24,18 @@ loaded_gfpgan_model = None
 def gfpgan():

     global loaded_gfpgan_model

 

-    if loaded_gfpgan_model is None and gfpgan_constructor is not None:

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

+    if loaded_gfpgan_model is not None:

+        return loaded_gfpgan_model

 

-    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

+

+    return model

 

 

 def gfpgan_fix_faces(np_image):

diff --git a/modules/shared.py b/modules/shared.py
index b02c9e36..c8c2749a 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -25,6 +25,7 @@ parser.add_argument("--allow-code", action='store_true', help="allow custom scri
 parser.add_argument("--medvram", action='store_true', help="enable stable diffusion model optimizations for sacrficing a little speed for low VRM usage")

 parser.add_argument("--lowvram", action='store_true', help="enable stable diffusion model optimizations for sacrficing a lot of speed for very low VRM usage")

 parser.add_argument("--always-batch-cond-uncond", action='store_true', help="a workaround test; may help with speed in you use --lowvram")

+parser.add_argument("--unload-gfpgan", action='store_true', help="unload GFPGAN every time after processing images. Warning: seems to cause memory leaks")

 parser.add_argument("--precision", type=str, help="evaluate at this precision", choices=["full", "autocast"], default="autocast")

 parser.add_argument("--share", action='store_true', help="use share=True for gradio and make the UI accessible through their site (doesn't work for me but you might have better luck)")

 cmd_opts = parser.parse_args()