From 65ccd6305fcf72347d5ed68f03095dced865ef6e Mon Sep 17 00:00:00 2001
From: v0xie <28695009+v0xie@users.noreply.github.com>
Date: Thu, 2 Nov 2023 00:11:32 -0700
Subject: detect diag_oft type

---
 extensions-builtin/Lora/networks.py | 7 +++++++
 1 file changed, 7 insertions(+)

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

diff --git a/extensions-builtin/Lora/networks.py b/extensions-builtin/Lora/networks.py
index 78a97033..7f814706 100644
--- a/extensions-builtin/Lora/networks.py
+++ b/extensions-builtin/Lora/networks.py
@@ -191,10 +191,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
-- 
cgit v1.2.1