From 4cbbb881ee530d9b9ba18027e2b0057e6a2c4ee1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=CE=A6=CF=86?= <42910943+Brawlence@users.noreply.github.com>
Date: Thu, 9 Mar 2023 07:56:19 +0300
Subject: Unload checkpoints on Request
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

…to free VRAM.

New Action buttons in the settings to manually free and reload checkpoints, essentially
juggling models between RAM and VRAM.
---
 modules/sd_models.py | 22 +++++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

(limited to 'modules/sd_models.py')

diff --git a/modules/sd_models.py b/modules/sd_models.py
index f0cb1240..f9dd0521 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -494,7 +494,7 @@ def reload_model_weights(sd_model=None, info=None):
     if sd_model is None or checkpoint_config != sd_model.used_config:
         del sd_model
         checkpoints_loaded.clear()
-        load_model(checkpoint_info, already_loaded_state_dict=state_dict, time_taken_to_load_state_dict=timer.records["load weights from disk"])
+        load_model(checkpoint_info, already_loaded_state_dict=state_dict)
         return shared.sd_model
 
     try:
@@ -517,3 +517,23 @@ def reload_model_weights(sd_model=None, info=None):
     print(f"Weights loaded in {timer.summary()}.")
 
     return sd_model
+
+def unload_model_weights(sd_model=None, info=None):
+    from modules import lowvram, devices, sd_hijack
+    timer = Timer()
+
+    if shared.sd_model:
+
+        # shared.sd_model.cond_stage_model.to(devices.cpu)
+        # shared.sd_model.first_stage_model.to(devices.cpu)
+        shared.sd_model.to(devices.cpu)
+        sd_hijack.model_hijack.undo_hijack(shared.sd_model)
+        shared.sd_model = None
+        sd_model = None
+        gc.collect()
+        devices.torch_gc()
+        torch.cuda.empty_cache()
+
+    print(f"Unloaded weights {timer.summary()}.")
+
+    return sd_model
\ No newline at end of file
-- 
cgit v1.2.3


From 92e173d414e5441ff34e5a6b12d5a766e86b0c36 Mon Sep 17 00:00:00 2001
From: carat-johyun <johyun@carat.im>
Date: Thu, 23 Mar 2023 14:28:08 +0900
Subject: fix variable typo

---
 modules/sd_models.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

(limited to 'modules/sd_models.py')

diff --git a/modules/sd_models.py b/modules/sd_models.py
index f0cb1240..6410b09a 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -178,7 +178,7 @@ def select_checkpoint():
     return checkpoint_info
 
 
-chckpoint_dict_replacements = {
+checkpoint_dict_replacements = {
     'cond_stage_model.transformer.embeddings.': 'cond_stage_model.transformer.text_model.embeddings.',
     'cond_stage_model.transformer.encoder.': 'cond_stage_model.transformer.text_model.encoder.',
     'cond_stage_model.transformer.final_layer_norm.': 'cond_stage_model.transformer.text_model.final_layer_norm.',
@@ -186,7 +186,7 @@ chckpoint_dict_replacements = {
 
 
 def transform_checkpoint_dict_key(k):
-    for text, replacement in chckpoint_dict_replacements.items():
+    for text, replacement in checkpoint_dict_replacements.items():
         if k.startswith(text):
             k = replacement + k[len(text):]
 
-- 
cgit v1.2.3