1 files changed, 22 insertions, 15 deletions
diff --git a/extensions-builtin/Lora/networks.py b/extensions-builtin/Lora/networks.py
index 22fdff4a..96f935b2 100644
--- a/extensions-builtin/Lora/networks.py
+++ b/extensions-builtin/Lora/networks.py
@@ -2,6 +2,7 @@ import logging
 import os
 import re
 
+import lora_patches
 import network
 import network_lora
 import network_hada
@@ -303,7 +304,10 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn
     wanted_names = tuple((x.name, x.te_multiplier, x.unet_multiplier, x.dyn_dim) for x in loaded_networks)
 
     weights_backup = getattr(self, "network_weights_backup", None)
-    if weights_backup is None:
+    if weights_backup is None and wanted_names != ():
+        if current_names != ():
+            raise RuntimeError("no backup weights found and current weights are not unchanged")
+
         if isinstance(self, torch.nn.MultiheadAttention):
             weights_backup = (self.in_proj_weight.to(devices.cpu, copy=True), self.out_proj.weight.to(devices.cpu, copy=True))
         else:
@@ -418,74 +422,74 @@ def network_reset_cached_weight(self: Union[torch.nn.Conv2d, torch.nn.Linear]):
 
 def network_Linear_forward(self, input):
     if shared.opts.lora_functional:
-        return network_forward(self, input, torch.nn.Linear_forward_before_network)
+        return network_forward(self, input, originals.Linear_forward)
 
     network_apply_weights(self)
 
-    return torch.nn.Linear_forward_before_network(self, input)
+    return originals.Linear_forward(self, input)
 
 
 def network_Linear_load_state_dict(self, *args, **kwargs):
     network_reset_cached_weight(self)
 
-    return torch.nn.Linear_load_state_dict_before_network(self, *args, **kwargs)
+    return originals.Linear_load_state_dict(self, *args, **kwargs)
 
 
 def network_Conv2d_forward(self, input):
     if shared.opts.lora_functional:
-        return network_forward(self, input, torch.nn.Conv2d_forward_before_network)
+        return network_forward(self, input, originals.Conv2d_forward)
 
     network_apply_weights(self)
 
-    return torch.nn.Conv2d_forward_before_network(self, input)
+    return originals.Conv2d_forward(self, input)
 
 
 def network_Conv2d_load_state_dict(self, *args, **kwargs):
     network_reset_cached_weight(self)
 
-    return torch.nn.Conv2d_load_state_dict_before_network(self, *args, **kwargs)
+    return originals.Conv2d_load_state_dict(self, *args, **kwargs)
 
 
 def network_GroupNorm_forward(self, input):
     if shared.opts.lora_functional:
-        return network_forward(self, input, torch.nn.GroupNorm_forward_before_network)
+        return network_forward(self, input, originals.GroupNorm_forward)
 
     network_apply_weights(self)
 
-    return torch.nn.GroupNorm_forward_before_network(self, input)
+    return originals.GroupNorm_forward(self, input)
 
 
 def network_GroupNorm_load_state_dict(self, *args, **kwargs):
     network_reset_cached_weight(self)
 
-    return torch.nn.GroupNorm_load_state_dict_before_network(self, *args, **kwargs)
+    return originals.GroupNorm_load_state_dict(self, *args, **kwargs)
 
 
 def network_LayerNorm_forward(self, input):
     if shared.opts.lora_functional:
-        return network_forward(self, input, torch.nn.LayerNorm_forward_before_network)
+        return network_forward(self, input, originals.LayerNorm_forward)
 
     network_apply_weights(self)
 
-    return torch.nn.LayerNorm_forward_before_network(self, input)
+    return originals.LayerNorm_forward(self, input)
 
 
 def network_LayerNorm_load_state_dict(self, *args, **kwargs):
     network_reset_cached_weight(self)
 
-    return torch.nn.LayerNorm_load_state_dict_before_network(self, *args, **kwargs)
+    return originals.LayerNorm_load_state_dict(self, *args, **kwargs)
 
 
 def network_MultiheadAttention_forward(self, *args, **kwargs):
     network_apply_weights(self)
 
-    return torch.nn.MultiheadAttention_forward_before_network(self, *args, **kwargs)
+    return originals.MultiheadAttention_forward(self, *args, **kwargs)
 
 
 def network_MultiheadAttention_load_state_dict(self, *args, **kwargs):
     network_reset_cached_weight(self)
 
-    return torch.nn.MultiheadAttention_load_state_dict_before_network(self, *args, **kwargs)
+    return originals.MultiheadAttention_load_state_dict(self, *args, **kwargs)
 
 
 def list_available_networks():
@@ -552,6 +556,9 @@ def infotext_pasted(infotext, params):
     if added:
         params["Prompt"] += "\n" + "".join(added)
 
+
+originals: lora_patches.LoraPatches = None
+
 extra_network_lora = None
 
 available_networks = {}