1 files changed, 9 insertions, 6 deletions
diff --git a/extensions-builtin/Lora/scripts/lora_script.py b/extensions-builtin/Lora/scripts/lora_script.py
index a385ae94..dc329e81 100644
--- a/extensions-builtin/Lora/scripts/lora_script.py
+++ b/extensions-builtin/Lora/scripts/lora_script.py
@@ -9,8 +9,9 @@ from modules import script_callbacks, ui_extra_networks, extra_networks, shared
 
 def unload():
     torch.nn.Linear.forward = torch.nn.Linear_forward_before_lora
+    torch.nn.Linear._load_from_state_dict = torch.nn.Linear_load_state_dict_before_lora
     torch.nn.Conv2d.forward = torch.nn.Conv2d_forward_before_lora
-    torch.nn.modules.linear.NonDynamicallyQuantizableLinear.forward = torch.nn.NonDynamicallyQuantizableLinear_forward_before_lora
+    torch.nn.Conv2d._load_from_state_dict = torch.nn.Conv2d_load_state_dict_before_lora
 
 
 def before_ui():
@@ -21,15 +22,19 @@ def before_ui():
 if not hasattr(torch.nn, 'Linear_forward_before_lora'):
     torch.nn.Linear_forward_before_lora = torch.nn.Linear.forward
 
+if not hasattr(torch.nn, 'Linear_load_state_dict_before_lora'):
+    torch.nn.Linear_load_state_dict_before_lora = torch.nn.Linear._load_from_state_dict
+
 if not hasattr(torch.nn, 'Conv2d_forward_before_lora'):
     torch.nn.Conv2d_forward_before_lora = torch.nn.Conv2d.forward
 
-if not hasattr(torch.nn, 'NonDynamicallyQuantizableLinear_forward_before_lora'):
-    torch.nn.NonDynamicallyQuantizableLinear_forward_before_lora = torch.nn.modules.linear.NonDynamicallyQuantizableLinear.forward
+if not hasattr(torch.nn, 'Conv2d_load_state_dict_before_lora'):
+    torch.nn.Conv2d_load_state_dict_before_lora = torch.nn.Conv2d._load_from_state_dict
 
 torch.nn.Linear.forward = lora.lora_Linear_forward
+torch.nn.Linear._load_from_state_dict = lora.lora_Linear_load_state_dict
 torch.nn.Conv2d.forward = lora.lora_Conv2d_forward
-torch.nn.modules.linear.NonDynamicallyQuantizableLinear.forward = lora.lora_NonDynamicallyQuantizableLinear_forward
+torch.nn.Conv2d._load_from_state_dict = lora.lora_Conv2d_load_state_dict
 
 script_callbacks.on_model_loaded(lora.assign_lora_names_to_compvis_modules)
 script_callbacks.on_script_unloaded(unload)
@@ -38,6 +43,4 @@ script_callbacks.on_before_ui(before_ui)
 
 shared.options_templates.update(shared.options_section(('extra_networks', "Extra Networks"), {
     "sd_lora": shared.OptionInfo("None", "Add Lora to prompt", gr.Dropdown, lambda: {"choices": [""] + [x for x in lora.available_loras]}, refresh=lora.list_available_loras),
-    "lora_apply_to_outputs": shared.OptionInfo(False, "Apply Lora to outputs rather than inputs when possible (experimental)"),
-
 }))