fix for loading of Kohya's Flux.2 dev lora

Author: christopher5106

src/diffusers/loaders/lora_conversion_utils.py

@@ -358,7 +358,10 @@ def _get_alpha_name(lora_name_alpha, diffusers_name, alpha):
 
 # The utilities under `_convert_kohya_flux_lora_to_diffusers()`
 # are adapted from https://github.com/kohya-ss/sd-scripts/blob/a61cf73a5cb5209c3f4d1a3688dd276a4dfd1ecb/networks/convert_flux_lora.py
-def _convert_kohya_flux_lora_to_diffusers(state_dict):
+def _convert_kohya_flux_lora_to_diffusers(
+    state_dict, 
+    version_flux2 = False,
+):
     def _convert_to_ai_toolkit(sds_sd, ait_sd, sds_key, ait_key):
         if sds_key + ".lora_down.weight" not in sds_sd:
             return
@@ -449,7 +452,15 @@ def _convert_to_ai_toolkit_cat(sds_sd, ait_sd, sds_key, ait_keys, dims=None):
 
     def _convert_sd_scripts_to_ai_toolkit(sds_sd):
         ait_sd = {}
-        for i in range(19):
+        
+        max_num_double_blocks, max_num_single_blocks = -1, -1
+        for key in list(sds_sd.keys()):
+            if key.startswith("lora_unet_double_blocks_"):
+                max_num_double_blocks = max(max_num_double_blocks, int(key.split("_")[4]))
+            if key.startswith("lora_unet_single_blocks_"):
+                max_num_single_blocks = max(max_num_single_blocks, int(key.split("_")[4]))
+
+        for i in range(max_num_double_blocks+1):
             _convert_to_ai_toolkit(
                 sds_sd,
                 ait_sd,
@@ -470,13 +481,21 @@ def _convert_sd_scripts_to_ai_toolkit(sds_sd):
                 sds_sd,
                 ait_sd,
                 f"lora_unet_double_blocks_{i}_img_mlp_0",
-                f"transformer.transformer_blocks.{i}.ff.net.0.proj",
+                (
+                    f"transformer.transformer_blocks.{i}.ff.linear_in"
+                    if version_flux2 else
+                    f"transformer.transformer_blocks.{i}.ff.net.0.proj"
+                ),
             )
             _convert_to_ai_toolkit(
                 sds_sd,
                 ait_sd,
                 f"lora_unet_double_blocks_{i}_img_mlp_2",
-                f"transformer.transformer_blocks.{i}.ff.net.2",
+                (
+                    f"transformer.transformer_blocks.{i}.ff.linear_out"
+                    if version_flux2 else
+                    f"transformer.transformer_blocks.{i}.ff.net.2"
+                ),
             )
             _convert_to_ai_toolkit(
                 sds_sd,
@@ -504,13 +523,21 @@ def _convert_sd_scripts_to_ai_toolkit(sds_sd):
                 sds_sd,
                 ait_sd,
                 f"lora_unet_double_blocks_{i}_txt_mlp_0",
-                f"transformer.transformer_blocks.{i}.ff_context.net.0.proj",
+                (
+                    f"transformer.transformer_blocks.{i}.ff_context.linear_in"
+                    if version_flux2 else
+                    f"transformer.transformer_blocks.{i}.ff_context.net.0.proj"
+                ),
             )
             _convert_to_ai_toolkit(
                 sds_sd,
                 ait_sd,
                 f"lora_unet_double_blocks_{i}_txt_mlp_2",
-                f"transformer.transformer_blocks.{i}.ff_context.net.2",
+                (
+                    f"transformer.transformer_blocks.{i}.ff_context.linear_out"
+                    if version_flux2 else
+                    f"transformer.transformer_blocks.{i}.ff_context.net.2"
+                ),
             )
             _convert_to_ai_toolkit(
                 sds_sd,
@@ -519,24 +546,36 @@ def _convert_sd_scripts_to_ai_toolkit(sds_sd):
                 f"transformer.transformer_blocks.{i}.norm1_context.linear",
             )
 
-        for i in range(38):
-            _convert_to_ai_toolkit_cat(
-                sds_sd,
-                ait_sd,
-                f"lora_unet_single_blocks_{i}_linear1",
-                [
-                    f"transformer.single_transformer_blocks.{i}.attn.to_q",
-                    f"transformer.single_transformer_blocks.{i}.attn.to_k",
-                    f"transformer.single_transformer_blocks.{i}.attn.to_v",
-                    f"transformer.single_transformer_blocks.{i}.proj_mlp",
-                ],
-                dims=[3072, 3072, 3072, 12288],
-            )
+        for i in range(max_num_single_blocks+1):
+            if version_flux2:
+                _convert_to_ai_toolkit(
+                    sds_sd,
+                    ait_sd,
+                    f"lora_unet_single_blocks_{i}_linear1",
+                    f"transformer.single_transformer_blocks.{i}.attn.to_qkv_mlp_proj",
+                )
+            else:
+                _convert_to_ai_toolkit_cat(
+                    sds_sd,
+                    ait_sd,
+                    f"lora_unet_single_blocks_{i}_linear1",
+                    [
+                        f"transformer.single_transformer_blocks.{i}.attn.to_q",
+                        f"transformer.single_transformer_blocks.{i}.attn.to_k",
+                        f"transformer.single_transformer_blocks.{i}.attn.to_v",
+                        f"transformer.single_transformer_blocks.{i}.proj_mlp",
+                    ],
+                    dims=[3072, 3072, 3072, 12288],
+                )
             _convert_to_ai_toolkit(
                 sds_sd,
                 ait_sd,
                 f"lora_unet_single_blocks_{i}_linear2",
-                f"transformer.single_transformer_blocks.{i}.proj_out",
+                (
+                    f"transformer.single_transformer_blocks.{i}.attn.to_out"
+                    if version_flux2 else
+                    f"transformer.single_transformer_blocks.{i}.proj_out"
+                ),
             )
             _convert_to_ai_toolkit(
                 sds_sd,

src/diffusers/loaders/lora_pipeline.py

@@ -5472,6 +5472,16 @@ def lora_state_dict(
             logger.warning(warn_msg)
             state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k}
 
+        is_kohya = any(".lora_down.weight" in k for k in state_dict)
+        if is_kohya:
+            state_dict = _convert_kohya_flux_lora_to_diffusers(
+                state_dict, 
+                version_flux2=True,
+            )
+            # Kohya already takes care of scaling the LoRA parameters with alpha.
+            for k in  state_dict:
+                assert "alpha" not in k, f"Found key with alpha: {k}"
+
         is_ai_toolkit = any(k.startswith("diffusion_model.") for k in state_dict)
         if is_ai_toolkit:
             state_dict = _convert_non_diffusers_flux2_lora_to_diffusers(state_dict)