1 files changed, 22 insertions, 3 deletions
diff --git a/extensions-builtin/Lora/scripts/lora_script.py b/extensions-builtin/Lora/scripts/lora_script.py
index cd28afc9..4c6e774a 100644
--- a/extensions-builtin/Lora/scripts/lora_script.py
+++ b/extensions-builtin/Lora/scripts/lora_script.py
@@ -23,9 +23,9 @@ def unload():
 def before_ui():
     ui_extra_networks.register_page(ui_extra_networks_lora.ExtraNetworksPageLora())
 
-    extra_network = extra_networks_lora.ExtraNetworkLora()
-    extra_networks.register_extra_network(extra_network)
-    extra_networks.register_extra_network_alias(extra_network, "lyco")
+    networks.extra_network_lora = extra_networks_lora.ExtraNetworkLora()
+    extra_networks.register_extra_network(networks.extra_network_lora)
+    extra_networks.register_extra_network_alias(networks.extra_network_lora, "lyco")
 
 
 if not hasattr(torch.nn, 'Linear_forward_before_network'):
@@ -40,6 +40,18 @@ if not hasattr(torch.nn, 'Conv2d_forward_before_network'):
 if not hasattr(torch.nn, 'Conv2d_load_state_dict_before_network'):
     torch.nn.Conv2d_load_state_dict_before_network = torch.nn.Conv2d._load_from_state_dict
 
+if not hasattr(torch.nn, 'GroupNorm_forward_before_network'):
+    torch.nn.GroupNorm_forward_before_network = torch.nn.GroupNorm.forward
+
+if not hasattr(torch.nn, 'GroupNorm_load_state_dict_before_network'):
+    torch.nn.GroupNorm_load_state_dict_before_network = torch.nn.GroupNorm._load_from_state_dict
+
+if not hasattr(torch.nn, 'LayerNorm_forward_before_network'):
+    torch.nn.LayerNorm_forward_before_network = torch.nn.LayerNorm.forward
+
+if not hasattr(torch.nn, 'LayerNorm_load_state_dict_before_network'):
+    torch.nn.LayerNorm_load_state_dict_before_network = torch.nn.LayerNorm._load_from_state_dict
+
 if not hasattr(torch.nn, 'MultiheadAttention_forward_before_network'):
     torch.nn.MultiheadAttention_forward_before_network = torch.nn.MultiheadAttention.forward
 
@@ -50,6 +62,10 @@ torch.nn.Linear.forward = networks.network_Linear_forward
 torch.nn.Linear._load_from_state_dict = networks.network_Linear_load_state_dict
 torch.nn.Conv2d.forward = networks.network_Conv2d_forward
 torch.nn.Conv2d._load_from_state_dict = networks.network_Conv2d_load_state_dict
+torch.nn.GroupNorm.forward = networks.network_GroupNorm_forward
+torch.nn.GroupNorm._load_from_state_dict = networks.network_GroupNorm_load_state_dict
+torch.nn.LayerNorm.forward = networks.network_LayerNorm_forward
+torch.nn.LayerNorm._load_from_state_dict = networks.network_LayerNorm_load_state_dict
 torch.nn.MultiheadAttention.forward = networks.network_MultiheadAttention_forward
 torch.nn.MultiheadAttention._load_from_state_dict = networks.network_MultiheadAttention_load_state_dict
 
@@ -65,6 +81,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 +138,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)