From b6bf15830acdb9bddf3d67698cc9a94ce9e2ac34 Mon Sep 17 00:00:00 2001 From: linoytsaban Date: Fri, 27 Jun 2025 17:40:38 +0300 Subject: [PATCH 1/7] initial commit --- .../loaders/lora_conversion_utils.py | 200 ++++++++++++++++++ src/diffusers/loaders/lora_pipeline.py | 11 + 2 files changed, 211 insertions(+) diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py index 25e06c007fbc..b2cb693190fd 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/src/diffusers/loaders/lora_conversion_utils.py @@ -1345,6 +1345,206 @@ def _convert_bfl_flux_control_lora_to_diffusers(original_state_dict): return converted_state_dict +def _convert_fal_kontext_lora_to_diffusers(original_state_dict): + converted_state_dict = {} + original_state_dict_keys = list(original_state_dict.keys()) + num_layers = 19 + num_single_layers = 38 + inner_dim = 3072 + mlp_ratio = 4.0 + + # double transformer blocks + for i in range(num_layers): + block_prefix = f"base_model.model." + + for lora_key in ["lora_A", "lora_B"]: + # norms + converted_state_dict[f"{block_prefix}norm1.linear.{lora_key}.weight"] = original_state_dict.pop( + f"double_blocks.{i}.img_mod.lin.{lora_key}.weight" + ) + if f"double_blocks.{i}.img_mod.lin.{lora_key}.bias" in original_state_dict_keys: + converted_state_dict[f"{block_prefix}norm1.linear.{lora_key}.bias"] = original_state_dict.pop( + f"double_blocks.{i}.img_mod.lin.{lora_key}.bias" + ) + + converted_state_dict[f"{block_prefix}norm1_context.linear.{lora_key}.weight"] = original_state_dict.pop( + f"double_blocks.{i}.txt_mod.lin.{lora_key}.weight" + ) + + # Q, K, V + if lora_key == "lora_A": + sample_lora_weight = original_state_dict.pop(f"double_blocks.{i}.img_attn.qkv.{lora_key}.weight") + converted_state_dict[f"{block_prefix}attn.to_v.{lora_key}.weight"] = torch.cat([sample_lora_weight]) + converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.weight"] = torch.cat([sample_lora_weight]) + converted_state_dict[f"{block_prefix}attn.to_k.{lora_key}.weight"] = torch.cat([sample_lora_weight]) + + context_lora_weight = original_state_dict.pop(f"double_blocks.{i}.txt_attn.qkv.{lora_key}.weight") + converted_state_dict[f"{block_prefix}attn.add_q_proj.{lora_key}.weight"] = torch.cat( + [context_lora_weight] + ) + converted_state_dict[f"{block_prefix}attn.add_k_proj.{lora_key}.weight"] = torch.cat( + [context_lora_weight] + ) + converted_state_dict[f"{block_prefix}attn.add_v_proj.{lora_key}.weight"] = torch.cat( + [context_lora_weight] + ) + else: + sample_q, sample_k, sample_v = torch.chunk( + original_state_dict.pop(f"double_blocks.{i}.img_attn.qkv.{lora_key}.weight"), 3, dim=0 + ) + converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.weight"] = torch.cat([sample_q]) + converted_state_dict[f"{block_prefix}attn.to_k.{lora_key}.weight"] = torch.cat([sample_k]) + converted_state_dict[f"{block_prefix}attn.to_v.{lora_key}.weight"] = torch.cat([sample_v]) + + context_q, context_k, context_v = torch.chunk( + original_state_dict.pop(f"double_blocks.{i}.txt_attn.qkv.{lora_key}.weight"), 3, dim=0 + ) + converted_state_dict[f"{block_prefix}attn.add_q_proj.{lora_key}.weight"] = torch.cat([context_q]) + converted_state_dict[f"{block_prefix}attn.add_k_proj.{lora_key}.weight"] = torch.cat([context_k]) + converted_state_dict[f"{block_prefix}attn.add_v_proj.{lora_key}.weight"] = torch.cat([context_v]) + + if f"double_blocks.{i}.img_attn.qkv.{lora_key}.bias" in original_state_dict_keys: + sample_q_bias, sample_k_bias, sample_v_bias = torch.chunk( + original_state_dict.pop(f"double_blocks.{i}.img_attn.qkv.{lora_key}.bias"), 3, dim=0 + ) + converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.bias"] = torch.cat([sample_q_bias]) + converted_state_dict[f"{block_prefix}attn.to_k.{lora_key}.bias"] = torch.cat([sample_k_bias]) + converted_state_dict[f"{block_prefix}attn.to_v.{lora_key}.bias"] = torch.cat([sample_v_bias]) + + if f"double_blocks.{i}.txt_attn.qkv.{lora_key}.bias" in original_state_dict_keys: + context_q_bias, context_k_bias, context_v_bias = torch.chunk( + original_state_dict.pop(f"double_blocks.{i}.txt_attn.qkv.{lora_key}.bias"), 3, dim=0 + ) + converted_state_dict[f"{block_prefix}attn.add_q_proj.{lora_key}.bias"] = torch.cat([context_q_bias]) + converted_state_dict[f"{block_prefix}attn.add_k_proj.{lora_key}.bias"] = torch.cat([context_k_bias]) + converted_state_dict[f"{block_prefix}attn.add_v_proj.{lora_key}.bias"] = torch.cat([context_v_bias]) + + # ff img_mlp + converted_state_dict[f"{block_prefix}ff.net.0.proj.{lora_key}.weight"] = original_state_dict.pop( + f"double_blocks.{i}.img_mlp.0.{lora_key}.weight" + ) + if f"double_blocks.{i}.img_mlp.0.{lora_key}.bias" in original_state_dict_keys: + converted_state_dict[f"{block_prefix}ff.net.0.proj.{lora_key}.bias"] = original_state_dict.pop( + f"double_blocks.{i}.img_mlp.0.{lora_key}.bias" + ) + + converted_state_dict[f"{block_prefix}ff.net.2.{lora_key}.weight"] = original_state_dict.pop( + f"double_blocks.{i}.img_mlp.2.{lora_key}.weight" + ) + if f"double_blocks.{i}.img_mlp.2.{lora_key}.bias" in original_state_dict_keys: + converted_state_dict[f"{block_prefix}ff.net.2.{lora_key}.bias"] = original_state_dict.pop( + f"double_blocks.{i}.img_mlp.2.{lora_key}.bias" + ) + + converted_state_dict[f"{block_prefix}ff_context.net.0.proj.{lora_key}.weight"] = original_state_dict.pop( + f"double_blocks.{i}.txt_mlp.0.{lora_key}.weight" + ) + if f"double_blocks.{i}.txt_mlp.0.{lora_key}.bias" in original_state_dict_keys: + converted_state_dict[f"{block_prefix}ff_context.net.0.proj.{lora_key}.bias"] = original_state_dict.pop( + f"double_blocks.{i}.txt_mlp.0.{lora_key}.bias" + ) + + converted_state_dict[f"{block_prefix}ff_context.net.2.{lora_key}.weight"] = original_state_dict.pop( + f"double_blocks.{i}.txt_mlp.2.{lora_key}.weight" + ) + if f"double_blocks.{i}.txt_mlp.2.{lora_key}.bias" in original_state_dict_keys: + converted_state_dict[f"{block_prefix}ff_context.net.2.{lora_key}.bias"] = original_state_dict.pop( + f"double_blocks.{i}.txt_mlp.2.{lora_key}.bias" + ) + + # output projections. + converted_state_dict[f"{block_prefix}attn.to_out.0.{lora_key}.weight"] = original_state_dict.pop( + f"double_blocks.{i}.img_attn.proj.{lora_key}.weight" + ) + if f"double_blocks.{i}.img_attn.proj.{lora_key}.bias" in original_state_dict_keys: + converted_state_dict[f"{block_prefix}attn.to_out.0.{lora_key}.bias"] = original_state_dict.pop( + f"double_blocks.{i}.img_attn.proj.{lora_key}.bias" + ) + converted_state_dict[f"{block_prefix}attn.to_add_out.{lora_key}.weight"] = original_state_dict.pop( + f"double_blocks.{i}.txt_attn.proj.{lora_key}.weight" + ) + if f"double_blocks.{i}.txt_attn.proj.{lora_key}.bias" in original_state_dict_keys: + converted_state_dict[f"{block_prefix}attn.to_add_out.{lora_key}.bias"] = original_state_dict.pop( + f"double_blocks.{i}.txt_attn.proj.{lora_key}.bias" + ) + + + # single transformer blocks + for i in range(num_single_layers): + block_prefix = f"single_transformer_blocks.{i}." + + for lora_key in ["lora_A", "lora_B"]: + # norm.linear <- single_blocks.0.modulation.lin + converted_state_dict[f"{block_prefix}norm.linear.{lora_key}.weight"] = original_state_dict.pop( + f"single_blocks.{i}.modulation.lin.{lora_key}.weight" + ) + if f"single_blocks.{i}.modulation.lin.{lora_key}.bias" in original_state_dict_keys: + converted_state_dict[f"{block_prefix}norm.linear.{lora_key}.bias"] = original_state_dict.pop( + f"single_blocks.{i}.modulation.lin.{lora_key}.bias" + ) + + # Q, K, V, mlp + mlp_hidden_dim = int(inner_dim * mlp_ratio) + split_size = (inner_dim, inner_dim, inner_dim, mlp_hidden_dim) + + if lora_key == "lora_A": + lora_weight = original_state_dict.pop(f"single_blocks.{i}.linear1.{lora_key}.weight") + converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.weight"] = torch.cat([lora_weight]) + converted_state_dict[f"{block_prefix}attn.to_k.{lora_key}.weight"] = torch.cat([lora_weight]) + converted_state_dict[f"{block_prefix}attn.to_v.{lora_key}.weight"] = torch.cat([lora_weight]) + converted_state_dict[f"{block_prefix}proj_mlp.{lora_key}.weight"] = torch.cat([lora_weight]) + + if f"single_blocks.{i}.linear1.{lora_key}.bias" in original_state_dict_keys: + lora_bias = original_state_dict.pop(f"single_blocks.{i}.linear1.{lora_key}.bias") + converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.bias"] = torch.cat([lora_bias]) + converted_state_dict[f"{block_prefix}attn.to_k.{lora_key}.bias"] = torch.cat([lora_bias]) + converted_state_dict[f"{block_prefix}attn.to_v.{lora_key}.bias"] = torch.cat([lora_bias]) + converted_state_dict[f"{block_prefix}proj_mlp.{lora_key}.bias"] = torch.cat([lora_bias]) + else: + q, k, v, mlp = torch.split( + original_state_dict.pop(f"single_blocks.{i}.linear1.{lora_key}.weight"), split_size, dim=0 + ) + converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.weight"] = torch.cat([q]) + converted_state_dict[f"{block_prefix}attn.to_k.{lora_key}.weight"] = torch.cat([k]) + converted_state_dict[f"{block_prefix}attn.to_v.{lora_key}.weight"] = torch.cat([v]) + converted_state_dict[f"{block_prefix}proj_mlp.{lora_key}.weight"] = torch.cat([mlp]) + + if f"single_blocks.{i}.linear1.{lora_key}.bias" in original_state_dict_keys: + q_bias, k_bias, v_bias, mlp_bias = torch.split( + original_state_dict.pop(f"single_blocks.{i}.linear1.{lora_key}.bias"), split_size, dim=0 + ) + converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.bias"] = torch.cat([q_bias]) + converted_state_dict[f"{block_prefix}attn.to_k.{lora_key}.bias"] = torch.cat([k_bias]) + converted_state_dict[f"{block_prefix}attn.to_v.{lora_key}.bias"] = torch.cat([v_bias]) + converted_state_dict[f"{block_prefix}proj_mlp.{lora_key}.bias"] = torch.cat([mlp_bias]) + + # output projections. + converted_state_dict[f"{block_prefix}proj_out.{lora_key}.weight"] = original_state_dict.pop( + f"single_blocks.{i}.linear2.{lora_key}.weight" + ) + if f"single_blocks.{i}.linear2.{lora_key}.bias" in original_state_dict_keys: + converted_state_dict[f"{block_prefix}proj_out.{lora_key}.bias"] = original_state_dict.pop( + f"single_blocks.{i}.linear2.{lora_key}.bias" + ) + + + for lora_key in ["lora_A", "lora_B"]: + converted_state_dict[f"proj_out.{lora_key}.weight"] = original_state_dict.pop( + f"final_layer.linear.{lora_key}.weight" + ) + if f"final_layer.linear.{lora_key}.bias" in original_state_dict_keys: + converted_state_dict[f"proj_out.{lora_key}.bias"] = original_state_dict.pop( + f"final_layer.linear.{lora_key}.bias" + ) + + if len(original_state_dict) > 0: + raise ValueError(f"`original_state_dict` should be empty at this point but has {original_state_dict.keys()=}.") + + for key in list(converted_state_dict.keys()): + converted_state_dict[f"transformer.{key}"] = converted_state_dict.pop(key) + + return converted_state_dict + def _convert_hunyuan_video_lora_to_diffusers(original_state_dict): converted_state_dict = {k: original_state_dict.pop(k) for k in list(original_state_dict.keys())} diff --git a/src/diffusers/loaders/lora_pipeline.py b/src/diffusers/loaders/lora_pipeline.py index 4fea005cbc39..bb33610b77c6 100644 --- a/src/diffusers/loaders/lora_pipeline.py +++ b/src/diffusers/loaders/lora_pipeline.py @@ -2062,6 +2062,17 @@ def lora_state_dict( return_metadata=return_lora_metadata, ) + is_fal_kontext = any("base_model.model" in k for k in state_dict) + if is_fal_kontext: + state_dict = _convert_fal_kontext_lora_to_diffusers(state_dict) + return cls._prepare_outputs( + state_dict, + metadata=metadata, + alphas=None, + return_alphas=return_alphas, + return_metadata=return_lora_metadata, + ) + # For state dicts like # https://huggingface.co/TheLastBen/Jon_Snow_Flux_LoRA keys = list(state_dict.keys()) From 04e2af0549454a05b9abf0c90ec1f4d643035c8e Mon Sep 17 00:00:00 2001 From: linoytsaban Date: Fri, 27 Jun 2025 17:52:07 +0300 Subject: [PATCH 2/7] initial commit --- src/diffusers/loaders/lora_pipeline.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/diffusers/loaders/lora_pipeline.py b/src/diffusers/loaders/lora_pipeline.py index bb33610b77c6..01474e713138 100644 --- a/src/diffusers/loaders/lora_pipeline.py +++ b/src/diffusers/loaders/lora_pipeline.py @@ -2062,7 +2062,8 @@ def lora_state_dict( return_metadata=return_lora_metadata, ) - is_fal_kontext = any("base_model.model" in k for k in state_dict) + is_fal_kontext = any("base_model" in k for k in state_dict) + print("SANITY CHECK: is_fal_kontext", is_fal_kontext) if is_fal_kontext: state_dict = _convert_fal_kontext_lora_to_diffusers(state_dict) return cls._prepare_outputs( From 38018e441c1e23e2030eeafd46ead9b1b94a7f23 Mon Sep 17 00:00:00 2001 From: linoytsaban Date: Fri, 27 Jun 2025 17:56:04 +0300 Subject: [PATCH 3/7] initial commit --- src/diffusers/loaders/lora_pipeline.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/diffusers/loaders/lora_pipeline.py b/src/diffusers/loaders/lora_pipeline.py index 01474e713138..5381d4a05783 100644 --- a/src/diffusers/loaders/lora_pipeline.py +++ b/src/diffusers/loaders/lora_pipeline.py @@ -1999,6 +1999,7 @@ def lora_state_dict( use_safetensors = kwargs.pop("use_safetensors", None) return_lora_metadata = kwargs.pop("return_lora_metadata", False) + print("RESPECTFULLY WTF") allow_pickle = False if use_safetensors is None: use_safetensors = True From 91de87fe9070d36149e467c1f349177ef7861535 Mon Sep 17 00:00:00 2001 From: linoytsaban Date: Fri, 27 Jun 2025 17:58:13 +0300 Subject: [PATCH 4/7] fix import --- src/diffusers/loaders/lora_pipeline.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/loaders/lora_pipeline.py b/src/diffusers/loaders/lora_pipeline.py index 5381d4a05783..939aa15034a7 100644 --- a/src/diffusers/loaders/lora_pipeline.py +++ b/src/diffusers/loaders/lora_pipeline.py @@ -41,6 +41,7 @@ ) from .lora_conversion_utils import ( _convert_bfl_flux_control_lora_to_diffusers, + _convert_fal_kontext_lora_to_diffusers, _convert_hunyuan_video_lora_to_diffusers, _convert_kohya_flux_lora_to_diffusers, _convert_musubi_wan_lora_to_diffusers, @@ -1999,7 +2000,6 @@ def lora_state_dict( use_safetensors = kwargs.pop("use_safetensors", None) return_lora_metadata = kwargs.pop("return_lora_metadata", False) - print("RESPECTFULLY WTF") allow_pickle = False if use_safetensors is None: use_safetensors = True From c1ba7ea11876056ad1d25b633355f67a18a9c429 Mon Sep 17 00:00:00 2001 From: linoytsaban Date: Fri, 27 Jun 2025 18:06:36 +0300 Subject: [PATCH 5/7] fix prefix --- .../loaders/lora_conversion_utils.py | 85 ++++++++++--------- 1 file changed, 43 insertions(+), 42 deletions(-) diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py index b2cb693190fd..d9150fd0e2e4 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/src/diffusers/loaders/lora_conversion_utils.py @@ -1355,30 +1355,31 @@ def _convert_fal_kontext_lora_to_diffusers(original_state_dict): # double transformer blocks for i in range(num_layers): - block_prefix = f"base_model.model." + block_prefix = f"transformer_blocks.{i}." + original_block_prefix = "base_model.model." for lora_key in ["lora_A", "lora_B"]: # norms converted_state_dict[f"{block_prefix}norm1.linear.{lora_key}.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_mod.lin.{lora_key}.weight" + f"{original_block_prefix}double_blocks.{i}.img_mod.lin.{lora_key}.weight" ) if f"double_blocks.{i}.img_mod.lin.{lora_key}.bias" in original_state_dict_keys: converted_state_dict[f"{block_prefix}norm1.linear.{lora_key}.bias"] = original_state_dict.pop( - f"double_blocks.{i}.img_mod.lin.{lora_key}.bias" + f"{original_block_prefix}double_blocks.{i}.img_mod.lin.{lora_key}.bias" ) converted_state_dict[f"{block_prefix}norm1_context.linear.{lora_key}.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_mod.lin.{lora_key}.weight" + f"{original_block_prefix}double_blocks.{i}.txt_mod.lin.{lora_key}.weight" ) # Q, K, V if lora_key == "lora_A": - sample_lora_weight = original_state_dict.pop(f"double_blocks.{i}.img_attn.qkv.{lora_key}.weight") + sample_lora_weight = original_state_dict.pop(f"{original_block_prefix}double_blocks.{i}.img_attn.qkv.{lora_key}.weight") converted_state_dict[f"{block_prefix}attn.to_v.{lora_key}.weight"] = torch.cat([sample_lora_weight]) converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.weight"] = torch.cat([sample_lora_weight]) converted_state_dict[f"{block_prefix}attn.to_k.{lora_key}.weight"] = torch.cat([sample_lora_weight]) - context_lora_weight = original_state_dict.pop(f"double_blocks.{i}.txt_attn.qkv.{lora_key}.weight") + context_lora_weight = original_state_dict.pop(f"{original_block_prefix}double_blocks.{i}.txt_attn.qkv.{lora_key}.weight") converted_state_dict[f"{block_prefix}attn.add_q_proj.{lora_key}.weight"] = torch.cat( [context_lora_weight] ) @@ -1390,14 +1391,14 @@ def _convert_fal_kontext_lora_to_diffusers(original_state_dict): ) else: sample_q, sample_k, sample_v = torch.chunk( - original_state_dict.pop(f"double_blocks.{i}.img_attn.qkv.{lora_key}.weight"), 3, dim=0 + original_state_dict.pop(f"{original_block_prefix}double_blocks.{i}.img_attn.qkv.{lora_key}.weight"), 3, dim=0 ) converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.weight"] = torch.cat([sample_q]) converted_state_dict[f"{block_prefix}attn.to_k.{lora_key}.weight"] = torch.cat([sample_k]) converted_state_dict[f"{block_prefix}attn.to_v.{lora_key}.weight"] = torch.cat([sample_v]) context_q, context_k, context_v = torch.chunk( - original_state_dict.pop(f"double_blocks.{i}.txt_attn.qkv.{lora_key}.weight"), 3, dim=0 + original_state_dict.pop(f"{original_block_prefix}double_blocks.{i}.txt_attn.qkv.{lora_key}.weight"), 3, dim=0 ) converted_state_dict[f"{block_prefix}attn.add_q_proj.{lora_key}.weight"] = torch.cat([context_q]) converted_state_dict[f"{block_prefix}attn.add_k_proj.{lora_key}.weight"] = torch.cat([context_k]) @@ -1405,7 +1406,7 @@ def _convert_fal_kontext_lora_to_diffusers(original_state_dict): if f"double_blocks.{i}.img_attn.qkv.{lora_key}.bias" in original_state_dict_keys: sample_q_bias, sample_k_bias, sample_v_bias = torch.chunk( - original_state_dict.pop(f"double_blocks.{i}.img_attn.qkv.{lora_key}.bias"), 3, dim=0 + original_state_dict.pop(f"{original_block_prefix}double_blocks.{i}.img_attn.qkv.{lora_key}.bias"), 3, dim=0 ) converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.bias"] = torch.cat([sample_q_bias]) converted_state_dict[f"{block_prefix}attn.to_k.{lora_key}.bias"] = torch.cat([sample_k_bias]) @@ -1413,7 +1414,7 @@ def _convert_fal_kontext_lora_to_diffusers(original_state_dict): if f"double_blocks.{i}.txt_attn.qkv.{lora_key}.bias" in original_state_dict_keys: context_q_bias, context_k_bias, context_v_bias = torch.chunk( - original_state_dict.pop(f"double_blocks.{i}.txt_attn.qkv.{lora_key}.bias"), 3, dim=0 + original_state_dict.pop(f"{original_block_prefix}double_blocks.{i}.txt_attn.qkv.{lora_key}.bias"), 3, dim=0 ) converted_state_dict[f"{block_prefix}attn.add_q_proj.{lora_key}.bias"] = torch.cat([context_q_bias]) converted_state_dict[f"{block_prefix}attn.add_k_proj.{lora_key}.bias"] = torch.cat([context_k_bias]) @@ -1421,51 +1422,51 @@ def _convert_fal_kontext_lora_to_diffusers(original_state_dict): # ff img_mlp converted_state_dict[f"{block_prefix}ff.net.0.proj.{lora_key}.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_mlp.0.{lora_key}.weight" + f"{original_block_prefix}double_blocks.{i}.img_mlp.0.{lora_key}.weight" ) - if f"double_blocks.{i}.img_mlp.0.{lora_key}.bias" in original_state_dict_keys: + if f"{original_block_prefix}double_blocks.{i}.img_mlp.0.{lora_key}.bias" in original_state_dict_keys: converted_state_dict[f"{block_prefix}ff.net.0.proj.{lora_key}.bias"] = original_state_dict.pop( - f"double_blocks.{i}.img_mlp.0.{lora_key}.bias" + f"{original_block_prefix}double_blocks.{i}.img_mlp.0.{lora_key}.bias" ) converted_state_dict[f"{block_prefix}ff.net.2.{lora_key}.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_mlp.2.{lora_key}.weight" + f"{original_block_prefix}double_blocks.{i}.img_mlp.2.{lora_key}.weight" ) - if f"double_blocks.{i}.img_mlp.2.{lora_key}.bias" in original_state_dict_keys: + if f"{original_block_prefix}double_blocks.{i}.img_mlp.2.{lora_key}.bias" in original_state_dict_keys: converted_state_dict[f"{block_prefix}ff.net.2.{lora_key}.bias"] = original_state_dict.pop( - f"double_blocks.{i}.img_mlp.2.{lora_key}.bias" + f"{original_block_prefix}double_blocks.{i}.img_mlp.2.{lora_key}.bias" ) converted_state_dict[f"{block_prefix}ff_context.net.0.proj.{lora_key}.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_mlp.0.{lora_key}.weight" + f"{original_block_prefix}double_blocks.{i}.txt_mlp.0.{lora_key}.weight" ) - if f"double_blocks.{i}.txt_mlp.0.{lora_key}.bias" in original_state_dict_keys: + if f"{original_block_prefix}double_blocks.{i}.txt_mlp.0.{lora_key}.bias" in original_state_dict_keys: converted_state_dict[f"{block_prefix}ff_context.net.0.proj.{lora_key}.bias"] = original_state_dict.pop( - f"double_blocks.{i}.txt_mlp.0.{lora_key}.bias" + f"{original_block_prefix}double_blocks.{i}.txt_mlp.0.{lora_key}.bias" ) converted_state_dict[f"{block_prefix}ff_context.net.2.{lora_key}.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_mlp.2.{lora_key}.weight" + f"{original_block_prefix}double_blocks.{i}.txt_mlp.2.{lora_key}.weight" ) - if f"double_blocks.{i}.txt_mlp.2.{lora_key}.bias" in original_state_dict_keys: + if f"{original_block_prefix}double_blocks.{i}.txt_mlp.2.{lora_key}.bias" in original_state_dict_keys: converted_state_dict[f"{block_prefix}ff_context.net.2.{lora_key}.bias"] = original_state_dict.pop( - f"double_blocks.{i}.txt_mlp.2.{lora_key}.bias" + f"{original_block_prefix}double_blocks.{i}.txt_mlp.2.{lora_key}.bias" ) # output projections. converted_state_dict[f"{block_prefix}attn.to_out.0.{lora_key}.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_attn.proj.{lora_key}.weight" + f"{original_block_prefix}double_blocks.{i}.img_attn.proj.{lora_key}.weight" ) - if f"double_blocks.{i}.img_attn.proj.{lora_key}.bias" in original_state_dict_keys: + if f"{original_block_prefix}double_blocks.{i}.img_attn.proj.{lora_key}.bias" in original_state_dict_keys: converted_state_dict[f"{block_prefix}attn.to_out.0.{lora_key}.bias"] = original_state_dict.pop( - f"double_blocks.{i}.img_attn.proj.{lora_key}.bias" + f"{original_block_prefix}double_blocks.{i}.img_attn.proj.{lora_key}.bias" ) converted_state_dict[f"{block_prefix}attn.to_add_out.{lora_key}.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_attn.proj.{lora_key}.weight" + f"{original_block_prefix}double_blocks.{i}.txt_attn.proj.{lora_key}.weight" ) - if f"double_blocks.{i}.txt_attn.proj.{lora_key}.bias" in original_state_dict_keys: + if f"{original_block_prefix}double_blocks.{i}.txt_attn.proj.{lora_key}.bias" in original_state_dict_keys: converted_state_dict[f"{block_prefix}attn.to_add_out.{lora_key}.bias"] = original_state_dict.pop( - f"double_blocks.{i}.txt_attn.proj.{lora_key}.bias" + f"{original_block_prefix}double_blocks.{i}.txt_attn.proj.{lora_key}.bias" ) @@ -1476,11 +1477,11 @@ def _convert_fal_kontext_lora_to_diffusers(original_state_dict): for lora_key in ["lora_A", "lora_B"]: # norm.linear <- single_blocks.0.modulation.lin converted_state_dict[f"{block_prefix}norm.linear.{lora_key}.weight"] = original_state_dict.pop( - f"single_blocks.{i}.modulation.lin.{lora_key}.weight" + f"{original_block_prefix}single_blocks.{i}.modulation.lin.{lora_key}.weight" ) - if f"single_blocks.{i}.modulation.lin.{lora_key}.bias" in original_state_dict_keys: + if f"{original_block_prefix}single_blocks.{i}.modulation.lin.{lora_key}.bias" in original_state_dict_keys: converted_state_dict[f"{block_prefix}norm.linear.{lora_key}.bias"] = original_state_dict.pop( - f"single_blocks.{i}.modulation.lin.{lora_key}.bias" + f"{original_block_prefix}single_blocks.{i}.modulation.lin.{lora_key}.bias" ) # Q, K, V, mlp @@ -1488,13 +1489,13 @@ def _convert_fal_kontext_lora_to_diffusers(original_state_dict): split_size = (inner_dim, inner_dim, inner_dim, mlp_hidden_dim) if lora_key == "lora_A": - lora_weight = original_state_dict.pop(f"single_blocks.{i}.linear1.{lora_key}.weight") + lora_weight = original_state_dict.pop(f"{original_block_prefix}single_blocks.{i}.linear1.{lora_key}.weight") converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.weight"] = torch.cat([lora_weight]) converted_state_dict[f"{block_prefix}attn.to_k.{lora_key}.weight"] = torch.cat([lora_weight]) converted_state_dict[f"{block_prefix}attn.to_v.{lora_key}.weight"] = torch.cat([lora_weight]) converted_state_dict[f"{block_prefix}proj_mlp.{lora_key}.weight"] = torch.cat([lora_weight]) - if f"single_blocks.{i}.linear1.{lora_key}.bias" in original_state_dict_keys: + if f"{original_block_prefix}single_blocks.{i}.linear1.{lora_key}.bias" in original_state_dict_keys: lora_bias = original_state_dict.pop(f"single_blocks.{i}.linear1.{lora_key}.bias") converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.bias"] = torch.cat([lora_bias]) converted_state_dict[f"{block_prefix}attn.to_k.{lora_key}.bias"] = torch.cat([lora_bias]) @@ -1502,16 +1503,16 @@ def _convert_fal_kontext_lora_to_diffusers(original_state_dict): converted_state_dict[f"{block_prefix}proj_mlp.{lora_key}.bias"] = torch.cat([lora_bias]) else: q, k, v, mlp = torch.split( - original_state_dict.pop(f"single_blocks.{i}.linear1.{lora_key}.weight"), split_size, dim=0 + original_state_dict.pop(f"{original_block_prefix}single_blocks.{i}.linear1.{lora_key}.weight"), split_size, dim=0 ) converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.weight"] = torch.cat([q]) converted_state_dict[f"{block_prefix}attn.to_k.{lora_key}.weight"] = torch.cat([k]) converted_state_dict[f"{block_prefix}attn.to_v.{lora_key}.weight"] = torch.cat([v]) converted_state_dict[f"{block_prefix}proj_mlp.{lora_key}.weight"] = torch.cat([mlp]) - if f"single_blocks.{i}.linear1.{lora_key}.bias" in original_state_dict_keys: + if f"{original_block_prefix}single_blocks.{i}.linear1.{lora_key}.bias" in original_state_dict_keys: q_bias, k_bias, v_bias, mlp_bias = torch.split( - original_state_dict.pop(f"single_blocks.{i}.linear1.{lora_key}.bias"), split_size, dim=0 + original_state_dict.pop(f"{original_block_prefix}single_blocks.{i}.linear1.{lora_key}.bias"), split_size, dim=0 ) converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.bias"] = torch.cat([q_bias]) converted_state_dict[f"{block_prefix}attn.to_k.{lora_key}.bias"] = torch.cat([k_bias]) @@ -1520,21 +1521,21 @@ def _convert_fal_kontext_lora_to_diffusers(original_state_dict): # output projections. converted_state_dict[f"{block_prefix}proj_out.{lora_key}.weight"] = original_state_dict.pop( - f"single_blocks.{i}.linear2.{lora_key}.weight" + f"{original_block_prefix}single_blocks.{i}.linear2.{lora_key}.weight" ) - if f"single_blocks.{i}.linear2.{lora_key}.bias" in original_state_dict_keys: + if f"{original_block_prefix}single_blocks.{i}.linear2.{lora_key}.bias" in original_state_dict_keys: converted_state_dict[f"{block_prefix}proj_out.{lora_key}.bias"] = original_state_dict.pop( - f"single_blocks.{i}.linear2.{lora_key}.bias" + f"{original_block_prefix}single_blocks.{i}.linear2.{lora_key}.bias" ) for lora_key in ["lora_A", "lora_B"]: converted_state_dict[f"proj_out.{lora_key}.weight"] = original_state_dict.pop( - f"final_layer.linear.{lora_key}.weight" + f"{original_block_prefix}final_layer.linear.{lora_key}.weight" ) - if f"final_layer.linear.{lora_key}.bias" in original_state_dict_keys: + if f"{original_block_prefix}final_layer.linear.{lora_key}.bias" in original_state_dict_keys: converted_state_dict[f"proj_out.{lora_key}.bias"] = original_state_dict.pop( - f"final_layer.linear.{lora_key}.bias" + f"{original_block_prefix}final_layer.linear.{lora_key}.bias" ) if len(original_state_dict) > 0: From ff743d03498dbdf034318cb9bb8e41597a1fd615 Mon Sep 17 00:00:00 2001 From: linoytsaban Date: Fri, 27 Jun 2025 18:13:23 +0300 Subject: [PATCH 6/7] remove print --- src/diffusers/loaders/lora_pipeline.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/diffusers/loaders/lora_pipeline.py b/src/diffusers/loaders/lora_pipeline.py index 939aa15034a7..4ee4808d801f 100644 --- a/src/diffusers/loaders/lora_pipeline.py +++ b/src/diffusers/loaders/lora_pipeline.py @@ -2064,7 +2064,6 @@ def lora_state_dict( ) is_fal_kontext = any("base_model" in k for k in state_dict) - print("SANITY CHECK: is_fal_kontext", is_fal_kontext) if is_fal_kontext: state_dict = _convert_fal_kontext_lora_to_diffusers(state_dict) return cls._prepare_outputs( From cb5a7fa968ea640a1c38bcb371dec794095904c3 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Fri, 27 Jun 2025 15:18:16 +0000 Subject: [PATCH 7/7] Apply style fixes --- .../loaders/lora_conversion_utils.py | 43 ++++++++++++++----- 1 file changed, 32 insertions(+), 11 deletions(-) diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py index d9150fd0e2e4..80929a1c8a0b 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/src/diffusers/loaders/lora_conversion_utils.py @@ -1345,6 +1345,7 @@ def _convert_bfl_flux_control_lora_to_diffusers(original_state_dict): return converted_state_dict + def _convert_fal_kontext_lora_to_diffusers(original_state_dict): converted_state_dict = {} original_state_dict_keys = list(original_state_dict.keys()) @@ -1374,12 +1375,16 @@ def _convert_fal_kontext_lora_to_diffusers(original_state_dict): # Q, K, V if lora_key == "lora_A": - sample_lora_weight = original_state_dict.pop(f"{original_block_prefix}double_blocks.{i}.img_attn.qkv.{lora_key}.weight") + sample_lora_weight = original_state_dict.pop( + f"{original_block_prefix}double_blocks.{i}.img_attn.qkv.{lora_key}.weight" + ) converted_state_dict[f"{block_prefix}attn.to_v.{lora_key}.weight"] = torch.cat([sample_lora_weight]) converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.weight"] = torch.cat([sample_lora_weight]) converted_state_dict[f"{block_prefix}attn.to_k.{lora_key}.weight"] = torch.cat([sample_lora_weight]) - context_lora_weight = original_state_dict.pop(f"{original_block_prefix}double_blocks.{i}.txt_attn.qkv.{lora_key}.weight") + context_lora_weight = original_state_dict.pop( + f"{original_block_prefix}double_blocks.{i}.txt_attn.qkv.{lora_key}.weight" + ) converted_state_dict[f"{block_prefix}attn.add_q_proj.{lora_key}.weight"] = torch.cat( [context_lora_weight] ) @@ -1391,14 +1396,22 @@ def _convert_fal_kontext_lora_to_diffusers(original_state_dict): ) else: sample_q, sample_k, sample_v = torch.chunk( - original_state_dict.pop(f"{original_block_prefix}double_blocks.{i}.img_attn.qkv.{lora_key}.weight"), 3, dim=0 + original_state_dict.pop( + f"{original_block_prefix}double_blocks.{i}.img_attn.qkv.{lora_key}.weight" + ), + 3, + dim=0, ) converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.weight"] = torch.cat([sample_q]) converted_state_dict[f"{block_prefix}attn.to_k.{lora_key}.weight"] = torch.cat([sample_k]) converted_state_dict[f"{block_prefix}attn.to_v.{lora_key}.weight"] = torch.cat([sample_v]) context_q, context_k, context_v = torch.chunk( - original_state_dict.pop(f"{original_block_prefix}double_blocks.{i}.txt_attn.qkv.{lora_key}.weight"), 3, dim=0 + original_state_dict.pop( + f"{original_block_prefix}double_blocks.{i}.txt_attn.qkv.{lora_key}.weight" + ), + 3, + dim=0, ) converted_state_dict[f"{block_prefix}attn.add_q_proj.{lora_key}.weight"] = torch.cat([context_q]) converted_state_dict[f"{block_prefix}attn.add_k_proj.{lora_key}.weight"] = torch.cat([context_k]) @@ -1406,7 +1419,9 @@ def _convert_fal_kontext_lora_to_diffusers(original_state_dict): if f"double_blocks.{i}.img_attn.qkv.{lora_key}.bias" in original_state_dict_keys: sample_q_bias, sample_k_bias, sample_v_bias = torch.chunk( - original_state_dict.pop(f"{original_block_prefix}double_blocks.{i}.img_attn.qkv.{lora_key}.bias"), 3, dim=0 + original_state_dict.pop(f"{original_block_prefix}double_blocks.{i}.img_attn.qkv.{lora_key}.bias"), + 3, + dim=0, ) converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.bias"] = torch.cat([sample_q_bias]) converted_state_dict[f"{block_prefix}attn.to_k.{lora_key}.bias"] = torch.cat([sample_k_bias]) @@ -1414,7 +1429,9 @@ def _convert_fal_kontext_lora_to_diffusers(original_state_dict): if f"double_blocks.{i}.txt_attn.qkv.{lora_key}.bias" in original_state_dict_keys: context_q_bias, context_k_bias, context_v_bias = torch.chunk( - original_state_dict.pop(f"{original_block_prefix}double_blocks.{i}.txt_attn.qkv.{lora_key}.bias"), 3, dim=0 + original_state_dict.pop(f"{original_block_prefix}double_blocks.{i}.txt_attn.qkv.{lora_key}.bias"), + 3, + dim=0, ) converted_state_dict[f"{block_prefix}attn.add_q_proj.{lora_key}.bias"] = torch.cat([context_q_bias]) converted_state_dict[f"{block_prefix}attn.add_k_proj.{lora_key}.bias"] = torch.cat([context_k_bias]) @@ -1469,7 +1486,6 @@ def _convert_fal_kontext_lora_to_diffusers(original_state_dict): f"{original_block_prefix}double_blocks.{i}.txt_attn.proj.{lora_key}.bias" ) - # single transformer blocks for i in range(num_single_layers): block_prefix = f"single_transformer_blocks.{i}." @@ -1489,7 +1505,9 @@ def _convert_fal_kontext_lora_to_diffusers(original_state_dict): split_size = (inner_dim, inner_dim, inner_dim, mlp_hidden_dim) if lora_key == "lora_A": - lora_weight = original_state_dict.pop(f"{original_block_prefix}single_blocks.{i}.linear1.{lora_key}.weight") + lora_weight = original_state_dict.pop( + f"{original_block_prefix}single_blocks.{i}.linear1.{lora_key}.weight" + ) converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.weight"] = torch.cat([lora_weight]) converted_state_dict[f"{block_prefix}attn.to_k.{lora_key}.weight"] = torch.cat([lora_weight]) converted_state_dict[f"{block_prefix}attn.to_v.{lora_key}.weight"] = torch.cat([lora_weight]) @@ -1503,7 +1521,9 @@ def _convert_fal_kontext_lora_to_diffusers(original_state_dict): converted_state_dict[f"{block_prefix}proj_mlp.{lora_key}.bias"] = torch.cat([lora_bias]) else: q, k, v, mlp = torch.split( - original_state_dict.pop(f"{original_block_prefix}single_blocks.{i}.linear1.{lora_key}.weight"), split_size, dim=0 + original_state_dict.pop(f"{original_block_prefix}single_blocks.{i}.linear1.{lora_key}.weight"), + split_size, + dim=0, ) converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.weight"] = torch.cat([q]) converted_state_dict[f"{block_prefix}attn.to_k.{lora_key}.weight"] = torch.cat([k]) @@ -1512,7 +1532,9 @@ def _convert_fal_kontext_lora_to_diffusers(original_state_dict): if f"{original_block_prefix}single_blocks.{i}.linear1.{lora_key}.bias" in original_state_dict_keys: q_bias, k_bias, v_bias, mlp_bias = torch.split( - original_state_dict.pop(f"{original_block_prefix}single_blocks.{i}.linear1.{lora_key}.bias"), split_size, dim=0 + original_state_dict.pop(f"{original_block_prefix}single_blocks.{i}.linear1.{lora_key}.bias"), + split_size, + dim=0, ) converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.bias"] = torch.cat([q_bias]) converted_state_dict[f"{block_prefix}attn.to_k.{lora_key}.bias"] = torch.cat([k_bias]) @@ -1528,7 +1550,6 @@ def _convert_fal_kontext_lora_to_diffusers(original_state_dict): f"{original_block_prefix}single_blocks.{i}.linear2.{lora_key}.bias" ) - for lora_key in ["lora_A", "lora_B"]: converted_state_dict[f"proj_out.{lora_key}.weight"] = original_state_dict.pop( f"{original_block_prefix}final_layer.linear.{lora_key}.weight"