From 22da90d4b8b1c5f582fb300b9b1e1d5808a02b60 Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Mon, 1 May 2023 10:13:21 -0400
Subject: [PATCH] fix lora memory leak

---
 TODO.md                         | 1 +
 extensions-builtin/Lora/lora.py | 3 +++
 modules/sd_models.py            | 2 ++
 modules/shared.py               | 4 ++--
 4 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/TODO.md b/TODO.md
index 040fca0a2..556370885 100644
--- a/TODO.md
+++ b/TODO.md
@@ -20,6 +20,7 @@ Stuff to be added...
 - Auto-test `torch.layer_norm` for FP16
 - Monitor file changes by misbehaving extensions
 - Kitchen theme: <https://github.com/canisminor1990/sd-webui-kitchen-theme>
+- Lightbox improvements
 
 ## Investigate
 
diff --git a/extensions-builtin/Lora/lora.py b/extensions-builtin/Lora/lora.py
index 3cbd91646..ac3f3a8e5 100644
--- a/extensions-builtin/Lora/lora.py
+++ b/extensions-builtin/Lora/lora.py
@@ -260,6 +260,9 @@ def lora_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.Mu
     If not, restores orginal weights from backup and alters weights according to loras.
     """
 
+    if len(loaded_loras) == 0:
+        return
+
     lora_layer_name = getattr(self, 'lora_layer_name', None)
     if lora_layer_name is None:
         return
diff --git a/modules/sd_models.py b/modules/sd_models.py
index a7a9adeb9..332a19e65 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -236,6 +236,7 @@ def read_metadata_from_safetensors(filename):
 
 def read_state_dict(checkpoint_file, map_location=None): # pylint: disable=unused-argument
     try:
+        pl_sd = None
         with progress.open(checkpoint_file, 'rb', description=f'Loading weights: [cyan]{checkpoint_file}', auto_refresh=True) as f:
             _, extension = os.path.splitext(checkpoint_file)
             if 'v1-5-pruned-emaonly.safetensors' or 'vae-ft-mse-840000-ema-pruned.ckpt' in checkpoint_file:
@@ -251,6 +252,7 @@ def read_state_dict(checkpoint_file, map_location=None): # pylint: disable=unuse
                     buffer = io.BytesIO(f.read())
                     pl_sd = torch.load(buffer, map_location='cpu')
             sd = get_state_dict_from_checkpoint(pl_sd)
+        del pl_sd
     except Exception as e:
         errors.display(e, f'loading model: {checkpoint_file}')
         sd = None
diff --git a/modules/shared.py b/modules/shared.py
index 611d8af4d..77d5b3e14 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -5,12 +5,12 @@ import json
 import datetime
 import gradio as gr
 import tqdm
+from modules import errors, ui_components, shared_items, cmd_args
+from modules.paths_internal import models_path, script_path, data_path, sd_configs_path, sd_default_config, sd_model_file, default_sd_model_file, extensions_dir, extensions_builtin_dir # pylint: disable=W0611
 import modules.interrogate
 import modules.memmon
 import modules.styles
 import modules.devices as devices
-from modules import errors, ui_components, shared_items, cmd_args
-from modules.paths_internal import models_path, script_path, data_path, sd_configs_path, sd_default_config, sd_model_file, default_sd_model_file, extensions_dir, extensions_builtin_dir # pylint: disable=W0611
 import modules.paths_internal as paths
 from installer import log as central_logger # pylint: disable=E0611