diff --git a/extensions-builtin/Lora/lora_convert.py b/extensions-builtin/Lora/lora_convert.py index 0c8497f85..04eb7bf84 100644 --- a/extensions-builtin/Lora/lora_convert.py +++ b/extensions-builtin/Lora/lora_convert.py @@ -161,7 +161,7 @@ class KeyConvert: sd_module.append(expert_module) i += 1 # SegMoe end - return key, sd_module + return [key], [sd_module] def diffusers(self, key): if self.is_sdxl: @@ -173,6 +173,27 @@ class KeyConvert: if search_key.startswith(map_key): key = key.replace(map_key, self.UNET_CONVERSION_MAP[map_key]).replace("oft", "lora") # pylint: disable=unsubscriptable-object sd_module = shared.sd_model.network_layer_mapping.get(key, None) + # SegMoE begin + expert_key = key + "_experts_0" + expert_module = shared.sd_model.network_layer_mapping.get(expert_key, None) + if expert_module is not None: + sd_module = expert_module + key = expert_key + if sd_module is None: + key = key.replace("_net_", "_experts_0_net_") + sd_module = shared.sd_model.network_layer_mapping.get(key, None) + key = key if isinstance(key, list) else [key] + sd_module = sd_module if isinstance(sd_module, list) else [sd_module] + if "_experts_0" in key[0]: + i = expert_module = 1 + while expert_module is not None: + expert_key = key[0].replace("_experts_0", f"_experts_{i}") + expert_module = shared.sd_model.network_layer_mapping.get(expert_key, None) + if expert_module is not None: + key.append(expert_key) + sd_module.append(expert_module) + i += 1 + # SegMoe end if debug and sd_module is None: raise RuntimeError(f"LoRA key not found in network_layer_mapping: key={key} mapping={shared.sd_model.network_layer_mapping.keys()}") return key, sd_module