fix: group aitoolkit lora layers

Author: Kevin Turner

invokeai/backend/patches/lora_conversions/flux_aitoolkit_lora_conversion_utils.py

@@ -1,12 +1,12 @@
 import json
-from collections import defaultdict
 from dataclasses import dataclass, field
 from typing import Any
 
 import torch
 
 from invokeai.backend.patches.layers.base_layer_patch import BaseLayerPatch
 from invokeai.backend.patches.layers.utils import any_lora_layer_from_state_dict
+from invokeai.backend.patches.lora_conversions.flux_diffusers_lora_conversion_utils import _group_by_layer
 from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_TRANSFORMER_PREFIX
 from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
 from invokeai.backend.util import InvokeAILogger
@@ -25,11 +25,11 @@ def is_state_dict_likely_in_flux_aitoolkit_format(state_dict: dict[str, Any], me
 
 @dataclass
 class GroupedStateDict:
-    transformer: dict = field(default_factory=dict)
+    transformer: dict[str, Any] = field(default_factory=dict)
     # might also grow CLIP and T5 submodels
 
 
-def _group_state_by_submodel(state_dict: dict[str, torch.Tensor]) -> GroupedStateDict:
+def _group_state_by_submodel(state_dict: dict[str, Any]) -> GroupedStateDict:
     logger = InvokeAILogger.get_logger()
     grouped = GroupedStateDict()
     for key, value in state_dict.items():
@@ -42,11 +42,22 @@ def _group_state_by_submodel(state_dict: dict[str, torch.Tensor]) -> GroupedStat
     return grouped
 
 
+def _rename_peft_lora_keys(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    """Renames keys from the PEFT LoRA format to the InvokeAI format."""
+    renamed_state_dict = {}
+    for key, value in state_dict.items():
+        renamed_key = key.replace(".lora_A.", ".lora_down.").replace(".lora_B.", ".lora_up.")
+        renamed_state_dict[renamed_key] = value
+    return renamed_state_dict
+
+
 def lora_model_from_flux_aitoolkit_state_dict(state_dict: dict[str, torch.Tensor]) -> ModelPatchRaw:
-    grouped = _group_state_by_submodel(state_dict)
+    state_dict = _rename_peft_lora_keys(state_dict)
+    by_layer = _group_by_layer(state_dict)
+    by_model = _group_state_by_submodel(by_layer)
 
     layers: dict[str, BaseLayerPatch] = {}
-    for layer_key, layer_state_dict in grouped.transformer.items():
+    for layer_key, layer_state_dict in by_model.transformer.items():
         layers[FLUX_LORA_TRANSFORMER_PREFIX + layer_key] = any_lora_layer_from_state_dict(layer_state_dict)
 
     return ModelPatchRaw(layers=layers)

tests/backend/patches/lora_conversions/test_flux_aitoolkit_lora_conversion_utils.py

@@ -11,10 +11,12 @@
 from tests.backend.patches.lora_conversions.lora_state_dicts.flux_dora_onetrainer_format import (
     state_dict_keys as flux_onetrainer_state_dict_keys,
 )
+from tests.backend.patches.lora_conversions.lora_state_dicts.flux_lora_aitoolkit_format import (
+    state_dict_keys as flux_aitoolkit_state_dict_keys,
+)
 from tests.backend.patches.lora_conversions.lora_state_dicts.flux_lora_diffusers_format import (
     state_dict_keys as flux_diffusers_state_dict_keys,
 )
-from tests.backend.patches.lora_conversions.lora_state_dicts.flux_lora_aitoolkit_format import state_dict_keys as flux_aitoolkit_state_dict_keys
 from tests.backend.patches.lora_conversions.lora_state_dicts.utils import keys_to_mock_state_dict
 
 
@@ -46,15 +48,12 @@ def test_flux_aitoolkit_transformer_state_dict_is_in_invoke_format():
     model_keys = set(model.state_dict().keys())
 
     for converted_key_prefix in converted_key_prefixes:
-        assert any(model_key.startswith(converted_key_prefix) for model_key in model_keys), f"'{converted_key_prefix}' did not match any model keys."
+        assert any(
+            model_key.startswith(converted_key_prefix) for model_key in model_keys
+        ), f"'{converted_key_prefix}' did not match any model keys."
 
 
 def test_lora_model_from_flux_aitoolkit_state_dict():
     state_dict = keys_to_mock_state_dict(flux_aitoolkit_state_dict_keys)
 
-    lora_model = lora_model_from_flux_aitoolkit_state_dict(state_dict)
-
-    # Assert that the lora_model has the expected layers.
-    # lora_model_keys = set(lora_model.layers.keys())
-    # lora_model_keys = {k.replace(".", "_") for k in lora_model_keys}
-    # assert lora_model_keys == expected_layer_keys
+    assert lora_model_from_flux_aitoolkit_state_dict(state_dict)