fix key convert for lora kohya lora (#183)

* fix key convert for lora from kohya

* simplify rename logic

Author: qzzz95

diffsynth_engine/pipelines/flux_image.py

@@ -45,7 +45,6 @@ def _from_kohya(self, lora_state_dict: Dict[str, torch.Tensor]) -> Dict[str, Dic
         flux_dim = 3072
         dit_rename_dict = flux_dit_config["civitai"]["rename_dict"]
         dit_suffix_rename_dict = flux_dit_config["civitai"]["suffix_rename_dict"]
-        clip_rename_dict = flux_text_encoder_config["diffusers"]["rename_dict"]
         clip_attn_rename_dict = flux_text_encoder_config["diffusers"]["attn_rename_dict"]
 
         dit_dict = {}
@@ -138,27 +137,18 @@ def _from_kohya(self, lora_state_dict: Dict[str, torch.Tensor]) -> Dict[str, Dic
                     lora_args["rank"] = lora_args["up"].shape[1]
                     rename = rename.replace(".weight", "")
                     dit_dict[rename] = lora_args
-            elif "lora_te" in key:
-                name = key.replace("lora_te1", "text_encoder")
-                name = name.replace("text_model_encoder_layers", "text_model.encoder.layers")
-                name = name.replace(".alpha", ".weight")
-                rename = ""
-                if name in clip_rename_dict:
-                    if name == "text_model.embeddings.position_embedding.weight":
-                        param = param.reshape((1, param.shape[0], param.shape[1]))
-                    rename = clip_rename_dict[name]
-                elif name.startswith("text_model.encoder.layers."):
-                    names = name.split(".")
-                    layer_id, layer_type, tail = names[3], ".".join(names[4:-1]), names[-1]
-                    rename = ".".join(["encoders", layer_id, clip_attn_rename_dict[layer_type], tail])
-                else:
-                    raise ValueError(f"Unsupported key: {key}")
+            elif "lora_te1_text_model_encoder_layers_" in key:
+                name = key.replace("lora_te1_text_model_encoder_layers_", "")
+                name = name.replace(".alpha", "")
+                layer_id, layer_type = name.split("_", 1)
+                layer_type = layer_type.replace("self_attn_", "self_attn.").replace("mlp_", "mlp.")
+                rename = ".".join(["encoders", layer_id, clip_attn_rename_dict[layer_type]])
+                
                 lora_args = {}
                 lora_args["alpha"] = param
                 lora_args["up"] = lora_state_dict[origin_key.replace(".alpha", ".lora_up.weight")]
                 lora_args["down"] = lora_state_dict[origin_key.replace(".alpha", ".lora_down.weight")]
                 lora_args["rank"] = lora_args["up"].shape[1]
-                rename = rename.replace(".weight", "")
                 te_dict[rename] = lora_args
             else:
                 raise ValueError(f"Unsupported key: {key}")