From 650ddc9dd3c1d126221682be8270f7fba1b5b6ce Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sun, 26 Mar 2023 10:44:20 +0300
Subject: Lora support for SD2

---
 extensions-builtin/Lora/scripts/lora_script.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

(limited to 'extensions-builtin/Lora/scripts/lora_script.py')

diff --git a/extensions-builtin/Lora/scripts/lora_script.py b/extensions-builtin/Lora/scripts/lora_script.py
index dc329e81..0adab225 100644
--- a/extensions-builtin/Lora/scripts/lora_script.py
+++ b/extensions-builtin/Lora/scripts/lora_script.py
@@ -12,6 +12,8 @@ def unload():
     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
+    torch.nn.MultiheadAttention.forward = torch.nn.MultiheadAttention_forward_before_lora
+    torch.nn.MultiheadAttention._load_from_state_dict = torch.nn.MultiheadAttention_load_state_dict_before_lora
 
 
 def before_ui():
@@ -31,10 +33,18 @@ if not hasattr(torch.nn, 'Conv2d_forward_before_lora'):
 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
 
+if not hasattr(torch.nn, 'MultiheadAttention_forward_before_lora'):
+    torch.nn.MultiheadAttention_forward_before_lora = torch.nn.MultiheadAttention.forward
+
+if not hasattr(torch.nn, 'MultiheadAttention_load_state_dict_before_lora'):
+    torch.nn.MultiheadAttention_load_state_dict_before_lora = torch.nn.MultiheadAttention._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
+torch.nn.MultiheadAttention.forward = lora.lora_MultiheadAttention_forward
+torch.nn.MultiheadAttention._load_from_state_dict = lora.lora_MultiheadAttention_load_state_dict
 
 script_callbacks.on_model_loaded(lora.assign_lora_names_to_compvis_modules)
 script_callbacks.on_script_unloaded(unload)
-- 
cgit v1.2.3