From 47666bb87079858c448ffe41a84dc71099b36c75 Mon Sep 17 00:00:00 2001 From: christopher5106 Date: Thu, 24 Sep 2026 13:54:25 +0200 Subject: [PATCH 1/2] [LoRA] accept fal-kontext LoRAs whose global embedder keys lack the base_model.model. prefix Some fal-kontext LoRAs store time_in / vector_in / txt_in / img_in / guidance_in without the `base_model.model.` prefix the block keys use, so `_convert_fal_kontext_lora_to_diffusers` left them in `original_state_dict` and raised "`original_state_dict` should be empty at this point". Map them to their diffusers names before that check. Rebuilt from scenario-labs/diffusers@02abf7eb1 (2026-04-15) without its Kohya Flux.2 hunks, which upstream now covers with `_convert_kohya_flux2_lora_to_diffusers`. --- .../loaders/lora_conversion_utils.py | 22 ++++++ tests/lora/test_lora_conversion_utils.py | 76 +++++++++++++++++++ 2 files changed, 98 insertions(+) create mode 100644 tests/lora/test_lora_conversion_utils.py diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py index 1b7bcc795d6b..bc111d77dd44 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/src/diffusers/loaders/lora_conversion_utils.py @@ -1582,6 +1582,28 @@ def _convert_fal_kontext_lora_to_diffusers(original_state_dict): f"{original_block_prefix}final_layer.linear.{lora_key}.bias" ) + # Some fal-kontext LoRAs carry the global embedder keys (time_in, vector_in, txt_in, img_in, guidance_in) + # without the `base_model.model.` prefix the block keys use. + for lora_key in ["lora_A", "lora_B"]: + for src, dst in [ + (f"time_in.in_layer.{lora_key}.weight", f"time_text_embed.timestep_embedder.linear_1.{lora_key}.weight"), + (f"time_in.out_layer.{lora_key}.weight", f"time_text_embed.timestep_embedder.linear_2.{lora_key}.weight"), + (f"vector_in.in_layer.{lora_key}.weight", f"time_text_embed.text_embedder.linear_1.{lora_key}.weight"), + (f"vector_in.out_layer.{lora_key}.weight", f"time_text_embed.text_embedder.linear_2.{lora_key}.weight"), + (f"txt_in.{lora_key}.weight", f"context_embedder.{lora_key}.weight"), + (f"img_in.{lora_key}.weight", f"x_embedder.{lora_key}.weight"), + ( + f"guidance_in.in_layer.{lora_key}.weight", + f"time_text_embed.guidance_embedder.linear_1.{lora_key}.weight", + ), + ( + f"guidance_in.out_layer.{lora_key}.weight", + f"time_text_embed.guidance_embedder.linear_2.{lora_key}.weight", + ), + ]: + if src in original_state_dict: + converted_state_dict[dst] = original_state_dict.pop(src) + if len(original_state_dict) > 0: raise ValueError(f"`original_state_dict` should be empty at this point but has {original_state_dict.keys()=}.") diff --git a/tests/lora/test_lora_conversion_utils.py b/tests/lora/test_lora_conversion_utils.py new file mode 100644 index 000000000000..2f287ad6c5f3 --- /dev/null +++ b/tests/lora/test_lora_conversion_utils.py @@ -0,0 +1,76 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import pytest +import torch + +from diffusers.loaders.lora_conversion_utils import _convert_fal_kontext_lora_to_diffusers + + +def _fal_kontext_state_dict(rank=1, inner_dim=3072, mlp_hidden_dim=12288, num_layers=19, num_single_layers=38): + """Minimal fal-kontext LoRA (block keys only), shaped so the qkv / linear1 splits work.""" + prefix = "base_model.model." + sd = {} + for i in range(num_layers): + for module in ["img_mod.lin", "txt_mod.lin", "img_mlp.0", "img_mlp.2", "txt_mlp.0", "txt_mlp.2"]: + sd[f"{prefix}double_blocks.{i}.{module}.lora_A.weight"] = torch.zeros(rank, 1) + sd[f"{prefix}double_blocks.{i}.{module}.lora_B.weight"] = torch.zeros(1, rank) + for module in ["img_attn.proj", "txt_attn.proj"]: + sd[f"{prefix}double_blocks.{i}.{module}.lora_A.weight"] = torch.zeros(rank, 1) + sd[f"{prefix}double_blocks.{i}.{module}.lora_B.weight"] = torch.zeros(1, rank) + for module in ["img_attn.qkv", "txt_attn.qkv"]: + sd[f"{prefix}double_blocks.{i}.{module}.lora_A.weight"] = torch.zeros(rank, 1) + sd[f"{prefix}double_blocks.{i}.{module}.lora_B.weight"] = torch.zeros(3 * inner_dim, rank) + for i in range(num_single_layers): + sd[f"{prefix}single_blocks.{i}.modulation.lin.lora_A.weight"] = torch.zeros(rank, 1) + sd[f"{prefix}single_blocks.{i}.modulation.lin.lora_B.weight"] = torch.zeros(1, rank) + sd[f"{prefix}single_blocks.{i}.linear1.lora_A.weight"] = torch.zeros(rank, 1) + sd[f"{prefix}single_blocks.{i}.linear1.lora_B.weight"] = torch.zeros(3 * inner_dim + mlp_hidden_dim, rank) + sd[f"{prefix}single_blocks.{i}.linear2.lora_A.weight"] = torch.zeros(rank, 1) + sd[f"{prefix}single_blocks.{i}.linear2.lora_B.weight"] = torch.zeros(1, rank) + sd[f"{prefix}final_layer.linear.lora_A.weight"] = torch.zeros(rank, 1) + sd[f"{prefix}final_layer.linear.lora_B.weight"] = torch.zeros(1, rank) + return sd + + +UNPREFIXED_GLOBAL_KEYS = { + "time_in.in_layer": "time_text_embed.timestep_embedder.linear_1", + "time_in.out_layer": "time_text_embed.timestep_embedder.linear_2", + "vector_in.in_layer": "time_text_embed.text_embedder.linear_1", + "vector_in.out_layer": "time_text_embed.text_embedder.linear_2", + "txt_in": "context_embedder", + "img_in": "x_embedder", + "guidance_in.in_layer": "time_text_embed.guidance_embedder.linear_1", + "guidance_in.out_layer": "time_text_embed.guidance_embedder.linear_2", +} + + +def test_fal_kontext_conversion_blocks_only(): + converted = _convert_fal_kontext_lora_to_diffusers(_fal_kontext_state_dict()) + assert all(k.startswith("transformer.") for k in converted) + assert "transformer.transformer_blocks.0.attn.to_q.lora_B.weight" in converted + assert "transformer.single_transformer_blocks.37.proj_mlp.lora_B.weight" in converted + + +@pytest.mark.parametrize("lora_key", ["lora_A", "lora_B"]) +def test_fal_kontext_conversion_accepts_unprefixed_global_embedders(lora_key): + sd = _fal_kontext_state_dict() + for src in UNPREFIXED_GLOBAL_KEYS: + sd[f"{src}.{lora_key}.weight"] = torch.ones(1, 1) + + converted = _convert_fal_kontext_lora_to_diffusers(sd) + + for src, dst in UNPREFIXED_GLOBAL_KEYS.items(): + assert f"transformer.{dst}.{lora_key}.weight" in converted, src + assert torch.equal(converted[f"transformer.{dst}.{lora_key}.weight"], torch.ones(1, 1)) From 2a27b3e2481a8c318ad33a8b681b08bd6ac879e5 Mon Sep 17 00:00:00 2001 From: christopher5106 Date: Tue, 29 Sep 2026 23:01:50 +0200 Subject: [PATCH 2/2] [LoRA] drop the unit test of the private fal-kontext converter diffusers does not unit-test private converter functions against third-party checkpoint layouts; the fix is reproducible through FluxPipeline.lora_state_dict (see the linked issue). --- tests/lora/test_lora_conversion_utils.py | 76 ------------------------ 1 file changed, 76 deletions(-) delete mode 100644 tests/lora/test_lora_conversion_utils.py diff --git a/tests/lora/test_lora_conversion_utils.py b/tests/lora/test_lora_conversion_utils.py deleted file mode 100644 index 2f287ad6c5f3..000000000000 --- a/tests/lora/test_lora_conversion_utils.py +++ /dev/null @@ -1,76 +0,0 @@ -# coding=utf-8 -# Copyright 2026 HuggingFace Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import pytest -import torch - -from diffusers.loaders.lora_conversion_utils import _convert_fal_kontext_lora_to_diffusers - - -def _fal_kontext_state_dict(rank=1, inner_dim=3072, mlp_hidden_dim=12288, num_layers=19, num_single_layers=38): - """Minimal fal-kontext LoRA (block keys only), shaped so the qkv / linear1 splits work.""" - prefix = "base_model.model." - sd = {} - for i in range(num_layers): - for module in ["img_mod.lin", "txt_mod.lin", "img_mlp.0", "img_mlp.2", "txt_mlp.0", "txt_mlp.2"]: - sd[f"{prefix}double_blocks.{i}.{module}.lora_A.weight"] = torch.zeros(rank, 1) - sd[f"{prefix}double_blocks.{i}.{module}.lora_B.weight"] = torch.zeros(1, rank) - for module in ["img_attn.proj", "txt_attn.proj"]: - sd[f"{prefix}double_blocks.{i}.{module}.lora_A.weight"] = torch.zeros(rank, 1) - sd[f"{prefix}double_blocks.{i}.{module}.lora_B.weight"] = torch.zeros(1, rank) - for module in ["img_attn.qkv", "txt_attn.qkv"]: - sd[f"{prefix}double_blocks.{i}.{module}.lora_A.weight"] = torch.zeros(rank, 1) - sd[f"{prefix}double_blocks.{i}.{module}.lora_B.weight"] = torch.zeros(3 * inner_dim, rank) - for i in range(num_single_layers): - sd[f"{prefix}single_blocks.{i}.modulation.lin.lora_A.weight"] = torch.zeros(rank, 1) - sd[f"{prefix}single_blocks.{i}.modulation.lin.lora_B.weight"] = torch.zeros(1, rank) - sd[f"{prefix}single_blocks.{i}.linear1.lora_A.weight"] = torch.zeros(rank, 1) - sd[f"{prefix}single_blocks.{i}.linear1.lora_B.weight"] = torch.zeros(3 * inner_dim + mlp_hidden_dim, rank) - sd[f"{prefix}single_blocks.{i}.linear2.lora_A.weight"] = torch.zeros(rank, 1) - sd[f"{prefix}single_blocks.{i}.linear2.lora_B.weight"] = torch.zeros(1, rank) - sd[f"{prefix}final_layer.linear.lora_A.weight"] = torch.zeros(rank, 1) - sd[f"{prefix}final_layer.linear.lora_B.weight"] = torch.zeros(1, rank) - return sd - - -UNPREFIXED_GLOBAL_KEYS = { - "time_in.in_layer": "time_text_embed.timestep_embedder.linear_1", - "time_in.out_layer": "time_text_embed.timestep_embedder.linear_2", - "vector_in.in_layer": "time_text_embed.text_embedder.linear_1", - "vector_in.out_layer": "time_text_embed.text_embedder.linear_2", - "txt_in": "context_embedder", - "img_in": "x_embedder", - "guidance_in.in_layer": "time_text_embed.guidance_embedder.linear_1", - "guidance_in.out_layer": "time_text_embed.guidance_embedder.linear_2", -} - - -def test_fal_kontext_conversion_blocks_only(): - converted = _convert_fal_kontext_lora_to_diffusers(_fal_kontext_state_dict()) - assert all(k.startswith("transformer.") for k in converted) - assert "transformer.transformer_blocks.0.attn.to_q.lora_B.weight" in converted - assert "transformer.single_transformer_blocks.37.proj_mlp.lora_B.weight" in converted - - -@pytest.mark.parametrize("lora_key", ["lora_A", "lora_B"]) -def test_fal_kontext_conversion_accepts_unprefixed_global_embedders(lora_key): - sd = _fal_kontext_state_dict() - for src in UNPREFIXED_GLOBAL_KEYS: - sd[f"{src}.{lora_key}.weight"] = torch.ones(1, 1) - - converted = _convert_fal_kontext_lora_to_diffusers(sd) - - for src, dst in UNPREFIXED_GLOBAL_KEYS.items(): - assert f"transformer.{dst}.{lora_key}.weight" in converted, src - assert torch.equal(converted[f"transformer.{dst}.{lora_key}.weight"], torch.ones(1, 1))