Merge branch 'AUTOMATIC1111:master' into master

1 files changed, 15 insertions, 1 deletions

diff --git a/extensions-builtin/Lora/networks.py b/extensions-builtin/Lora/networks.py
index 60d8dec4..629bf853 100644
--- a/extensions-builtin/Lora/networks.py
+++ b/extensions-builtin/Lora/networks.py
@@ -11,6 +11,7 @@ import network_ia3
 import network_lokr

 import network_full

 import network_norm

+import network_oft

 

 import torch

 from typing import Union

@@ -28,6 +29,7 @@ module_types = [
     network_full.ModuleTypeFull(),

     network_norm.ModuleTypeNorm(),

     network_glora.ModuleTypeGLora(),

+    network_oft.ModuleTypeOFT(),

 ]

 

 

@@ -157,7 +159,8 @@ def load_network(name, network_on_disk):
     bundle_embeddings = {}

 

     for key_network, weight in sd.items():

-        key_network_without_network_parts, network_part = key_network.split(".", 1)

+        key_network_without_network_parts, _, network_part = key_network.partition(".")

+

         if key_network_without_network_parts == "bundle_emb":

             emb_name, vec_name = network_part.split(".", 1)

             emb_dict = bundle_embeddings.get(emb_name, {})

@@ -189,6 +192,17 @@ def load_network(name, network_on_disk):
                 key = key_network_without_network_parts.replace("lora_te1_text_model", "transformer_text_model")

                 sd_module = shared.sd_model.network_layer_mapping.get(key, None)

 

+        # kohya_ss OFT module

+        elif sd_module is None and "oft_unet" in key_network_without_network_parts:

+            key = key_network_without_network_parts.replace("oft_unet", "diffusion_model")

+            sd_module = shared.sd_model.network_layer_mapping.get(key, None)

+

+        # KohakuBlueLeaf OFT module

+        if sd_module is None and "oft_diag" in key:

+            key = key_network_without_network_parts.replace("lora_unet", "diffusion_model")

+            key = key_network_without_network_parts.replace("lora_te1_text_model", "0_transformer_text_model")

+            sd_module = shared.sd_model.network_layer_mapping.get(key, None)

+

         if sd_module is None:

             keys_failed_to_match[key_network] = key

             continue