diff --git a/invokeai/backend/patches/lora_conversions/flux_diffusers_lora_conversion_utils.py b/invokeai/backend/patches/lora_conversions/flux_diffusers_lora_conversion_utils.py index d50f5d9cb5d..31e59409d68 100644 --- a/invokeai/backend/patches/lora_conversions/flux_diffusers_lora_conversion_utils.py +++ b/invokeai/backend/patches/lora_conversions/flux_diffusers_lora_conversion_utils.py @@ -5,7 +5,11 @@ from invokeai.backend.patches.layers.base_layer_patch import BaseLayerPatch from invokeai.backend.patches.layers.merged_layer_patch import MergedLayerPatch, Range from invokeai.backend.patches.layers.utils import any_lora_layer_from_state_dict -from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_TRANSFORMER_PREFIX +from invokeai.backend.patches.lora_conversions.flux_lora_constants import ( + FLUX_LORA_CLIP_PREFIX, + FLUX_LORA_T5_PREFIX, + FLUX_LORA_TRANSFORMER_PREFIX, +) from invokeai.backend.patches.model_patch_raw import ModelPatchRaw @@ -101,6 +105,22 @@ def lora_layers_from_flux_diffusers_grouped_state_dict( https://github.com/huggingface/diffusers/blob/55ac421f7bb12fd00ccbef727be4dc2f3f920abb/scripts/convert_flux_to_diffusers.py """ + # Diffusers keeps CLIP layers under "text_encoder." and T5 layers under "text_encoder_2.", + # alongside the "transformer." layers. Route them to the CLIP and T5 prefixes; only the + # transformer layers continue through the conversion below. + clip_grouped_sd: dict[str, dict[str, torch.Tensor]] = {} + t5_grouped_sd: dict[str, dict[str, torch.Tensor]] = {} + transformer_grouped_sd: dict[str, dict[str, torch.Tensor]] = {} + for k, v in grouped_state_dict.items(): + bare = k.removeprefix("base_model.model.") + if bare.startswith("text_encoder_2."): + t5_grouped_sd[bare.removeprefix("text_encoder_2.")] = v + elif bare.startswith("text_encoder."): + clip_grouped_sd[bare.removeprefix("text_encoder.")] = v + else: + transformer_grouped_sd[k] = v + grouped_state_dict = transformer_grouped_sd + # Determine which prefix is used and remove it from all keys. # Check if any key starts with "base_model.model." prefix has_base_model_prefix = any(k.startswith("base_model.model.") for k in grouped_state_dict.keys()) @@ -289,6 +309,14 @@ def add_qkv_lora_layer_if_present( layers_with_prefix = {f"{FLUX_LORA_TRANSFORMER_PREFIX}{k}": v for k, v in layers.items()} + for model_prefix, te_grouped_sd in [ + (FLUX_LORA_CLIP_PREFIX, clip_grouped_sd), + (FLUX_LORA_T5_PREFIX, t5_grouped_sd), + ]: + for layer_key, src_layer_dict in te_grouped_sd.items(): + values = get_lora_layer_values(src_layer_dict) + layers_with_prefix[f"{model_prefix}{layer_key}"] = any_lora_layer_from_state_dict(values) + return layers_with_prefix diff --git a/tests/backend/patches/lora_conversions/test_flux_diffusers_lora_conversion_utils.py b/tests/backend/patches/lora_conversions/test_flux_diffusers_lora_conversion_utils.py index 8bfdb82588b..2c1a9750df0 100644 --- a/tests/backend/patches/lora_conversions/test_flux_diffusers_lora_conversion_utils.py +++ b/tests/backend/patches/lora_conversions/test_flux_diffusers_lora_conversion_utils.py @@ -5,7 +5,11 @@ is_state_dict_likely_in_flux_diffusers_format, lora_model_from_flux_diffusers_state_dict, ) -from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_TRANSFORMER_PREFIX +from invokeai.backend.patches.lora_conversions.flux_lora_constants import ( + FLUX_LORA_CLIP_PREFIX, + FLUX_LORA_T5_PREFIX, + FLUX_LORA_TRANSFORMER_PREFIX, +) from tests.backend.patches.lora_conversions.lora_state_dicts.flux_dora_onetrainer_format import ( state_dict_keys as flux_onetrainer_state_dict_keys, ) @@ -117,3 +121,20 @@ def test_lora_model_from_flux_diffusers_state_dict_guidance_embedder(): out_layer = model.layers[f"{FLUX_LORA_TRANSFORMER_PREFIX}guidance_in.out_layer"] assert in_layer.get_parameters({"weight": torch.empty(3072, 256)}, 1.0)["weight"].shape == (3072, 256) assert out_layer.get_parameters({"weight": torch.empty(3072, 3072)}, 1.0)["weight"].shape == (3072, 3072) + + +def test_lora_model_from_flux_diffusers_state_dict_with_text_encoder_layers(): + """Test that a diffusers-format FLUX LoRA that also trains the text encoders loads, routing the + text_encoder.* keys to the CLIP prefix and text_encoder_2.* keys to the T5 prefix. + """ + state_dict = keys_to_mock_state_dict(flux_diffusers_state_dict_keys) + state_dict["text_encoder.text_model.encoder.layers.0.self_attn.q_proj.lora_A.weight"] = torch.empty([4, 768]) + state_dict["text_encoder.text_model.encoder.layers.0.self_attn.q_proj.lora_B.weight"] = torch.empty([768, 4]) + state_dict["text_encoder_2.encoder.block.0.layer.0.SelfAttention.q.lora_A.weight"] = torch.empty([4, 4096]) + state_dict["text_encoder_2.encoder.block.0.layer.0.SelfAttention.q.lora_B.weight"] = torch.empty([4096, 4]) + + model = lora_model_from_flux_diffusers_state_dict(state_dict, alpha=8.0) + + assert any(k.startswith(FLUX_LORA_TRANSFORMER_PREFIX) for k in model.layers) + assert f"{FLUX_LORA_CLIP_PREFIX}text_model.encoder.layers.0.self_attn.q_proj" in model.layers + assert f"{FLUX_LORA_T5_PREFIX}encoder.block.0.layer.0.SelfAttention.q" in model.layers