apply Lora by altering layer's weights instead of adding more calculations in forward()

1 files changed, 10 insertions, 2 deletions

diff --git a/extensions-builtin/Lora/scripts/lora_script.py b/extensions-builtin/Lora/scripts/lora_script.py
index 2e860160..dc329e81 100644
--- a/extensions-builtin/Lora/scripts/lora_script.py
+++ b/extensions-builtin/Lora/scripts/lora_script.py
@@ -9,7 +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.Conv2d._load_from_state_dict = torch.nn.Conv2d_load_state_dict_before_lora

 

 

 def before_ui():

@@ -20,11 +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, '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.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)

@@ -33,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)"),

-

 }))