From eed963e97261ee03bffe59e3d343dcf53d82dbfd Mon Sep 17 00:00:00 2001
From: AUTOMATIC1111 <16777216c@gmail.com>
Date: Wed, 9 Aug 2023 16:54:49 +0300
Subject: Lora cache in memory

---
 extensions-builtin/Lora/networks.py            | 22 +++++++++++++++++++---
 extensions-builtin/Lora/scripts/lora_script.py |  3 +++
 2 files changed, 22 insertions(+), 3 deletions(-)

(limited to 'extensions-builtin')

diff --git a/extensions-builtin/Lora/networks.py b/extensions-builtin/Lora/networks.py
index 17cbe1bb..bc722e90 100644
--- a/extensions-builtin/Lora/networks.py
+++ b/extensions-builtin/Lora/networks.py
@@ -195,6 +195,15 @@ def load_network(name, network_on_disk):
     return net
 
 
+def purge_networks_from_memory():
+    while len(networks_in_memory) > shared.opts.lora_in_memory_limit and len(networks_in_memory) > 0:
+        name = next(iter(networks_in_memory))
+        networks_in_memory.pop(name, None)
+
+    devices.torch_gc()
+
+
+
 def load_networks(names, te_multipliers=None, unet_multipliers=None, dyn_dims=None):
     already_loaded = {}
 
@@ -212,15 +221,19 @@ def load_networks(names, te_multipliers=None, unet_multipliers=None, dyn_dims=No
 
     failed_to_load_networks = []
 
-    for i, name in enumerate(names):
+    for i, (network_on_disk, name) in enumerate(zip(networks_on_disk, names)):
         net = already_loaded.get(name, None)
 
-        network_on_disk = networks_on_disk[i]
-
         if network_on_disk is not None:
+            if net is None:
+                net = networks_in_memory.get(name)
+
             if net is None or os.path.getmtime(network_on_disk.filename) > net.mtime:
                 try:
                     net = load_network(name, network_on_disk)
+
+                    networks_in_memory.pop(name, None)
+                    networks_in_memory[name] = net
                 except Exception as e:
                     errors.display(e, f"loading network {network_on_disk.filename}")
                     continue
@@ -242,6 +255,8 @@ def load_networks(names, te_multipliers=None, unet_multipliers=None, dyn_dims=No
     if failed_to_load_networks:
         sd_hijack.model_hijack.comments.append("Failed to find networks: " + ", ".join(failed_to_load_networks))
 
+    purge_networks_from_memory()
+
 
 def network_restore_weights_from_backup(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.MultiheadAttention]):
     weights_backup = getattr(self, "network_weights_backup", None)
@@ -462,6 +477,7 @@ def infotext_pasted(infotext, params):
 available_networks = {}
 available_network_aliases = {}
 loaded_networks = []
+networks_in_memory = {}
 available_network_hash_lookup = {}
 forbidden_network_aliases = {}
 
diff --git a/extensions-builtin/Lora/scripts/lora_script.py b/extensions-builtin/Lora/scripts/lora_script.py
index cd28afc9..6ab8b6e7 100644
--- a/extensions-builtin/Lora/scripts/lora_script.py
+++ b/extensions-builtin/Lora/scripts/lora_script.py
@@ -65,6 +65,7 @@ shared.options_templates.update(shared.options_section(('extra_networks', "Extra
     "lora_add_hashes_to_infotext": shared.OptionInfo(True, "Add Lora hashes to infotext"),
     "lora_show_all": shared.OptionInfo(False, "Always show all networks on the Lora page").info("otherwise, those detected as for incompatible version of Stable Diffusion will be hidden"),
     "lora_hide_unknown_for_versions": shared.OptionInfo([], "Hide networks of unknown versions for model versions", gr.CheckboxGroup, {"choices": ["SD1", "SD2", "SDXL"]}),
+    "lora_in_memory_limit": shared.OptionInfo(0, "Number of Lora networks to keep cached in memory", gr.Number, {"precision": 0}),
 }))
 
 
@@ -121,3 +122,5 @@ def infotext_pasted(infotext, d):
 
 
 script_callbacks.on_infotext_pasted(infotext_pasted)
+
+shared.opts.onchange("lora_in_memory_limit", networks.purge_networks_from_memory)
-- 
cgit v1.2.1