From 1f779226f0dc01d97fbc3ea556137fd19bd7e412 Mon Sep 17 00:00:00 2001 From: Won-Kyu Park Date: Thu, 19 Sep 2024 09:29:46 +0900 Subject: [PATCH] check lora_unet prefix to support Black Forest Labs's lora --- extensions-builtin/Lora/networks.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/extensions-builtin/Lora/networks.py b/extensions-builtin/Lora/networks.py index dfdf3c7e6..5f00d36e7 100644 --- a/extensions-builtin/Lora/networks.py +++ b/extensions-builtin/Lora/networks.py @@ -183,8 +183,12 @@ def load_network(name, network_on_disk): for key_network, weight in sd.items(): if diffusers_weight_map: - key_network_without_network_parts, network_name, network_weight = key_network.rsplit(".", 2) - network_part = network_name + '.' + network_weight + if key_network.startswith("lora_unet"): + key_network_without_network_parts, _, network_part = key_network.partition(".") + key_network_without_network_parts = key_network_without_network_parts.replace("lora_unet", "diffusion_model") + else: + key_network_without_network_parts, network_name, network_weight = key_network.rsplit(".", 2) + network_part = network_name + '.' + network_weight else: key_network_without_network_parts, _, network_part = key_network.partition(".")