From 19a9a78c9bb54d13a8f291f43939331af3242e8a Mon Sep 17 00:00:00 2001
From: DenOfEquity <166248528+DenOfEquity@users.noreply.github.com>
Date: Thu, 14 Nov 2024 22:05:54 +0000
Subject: [PATCH] fix/workaround for potential memory leak (#2315)

unload old models, based on reference count <= 2
(in practise only noticed extra copies of JointTextEncoder, never KModel or IntegratedAutoencoderKL)
#2281 #2308 and others
---
 backend/memory_management.py | 5 +++++
 modules/sd_models.py         | 2 +-
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/backend/memory_management.py b/backend/memory_management.py
index 66f89c07..5f0c8312 100644
--- a/backend/memory_management.py
+++ b/backend/memory_management.py
@@ -560,6 +560,11 @@ def unload_model_clones(model):
 
 
 def free_memory(memory_required, device, keep_loaded=[], free_all=False):
+    # this check fully unloads any 'abandoned' models
+    for i in range(len(current_loaded_models) - 1, -1, -1):
+        if sys.getrefcount(current_loaded_models[i].model) <= 2:
+            current_loaded_models.pop(i).model_unload(avoid_model_moving=True)
+
     if free_all:
         memory_required = 1e30
         print(f"[Unload] Trying to free all memory for {device} with {len(keep_loaded)} models keep loaded ... ", end="")
diff --git a/modules/sd_models.py b/modules/sd_models.py
index 3359603e..76940ecc 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -450,7 +450,7 @@ def reload_model_weights(sd_model=None, info=None, forced_reload=False):
 
 def unload_model_weights(sd_model=None, info=None):
     memory_management.unload_all_models()
-    pass
+    return
 
 
 def apply_token_merging(sd_model, token_merging_ratio):