fix(flux): load diffusers FLUX LoRAs with text-encoder layers

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
 
 

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