test: add some aitoolkit lora tests

Author: Kevin Turner

invokeai/backend/model_manager/load/model_loaders/lora.py

@@ -21,7 +21,7 @@
     SubModelType,
 )
 from invokeai.backend.patches.lora_conversions.flux_aitoolkit_lora_conversion_utils import (
-    is_state_dict_likely_in_aitoolkit_format,
+    is_state_dict_likely_in_flux_aitoolkit_format,
     lora_model_from_flux_aitoolkit_state_dict,
 )
 from invokeai.backend.patches.lora_conversions.flux_control_lora_utils import (
@@ -96,7 +96,7 @@ def _load_model(
                     model = lora_model_from_flux_onetrainer_state_dict(state_dict=state_dict)
                 elif is_state_dict_likely_flux_control(state_dict=state_dict):
                     model = lora_model_from_flux_control_state_dict(state_dict=state_dict)
-                elif is_state_dict_likely_in_aitoolkit_format(state_dict=state_dict):
+                elif is_state_dict_likely_in_flux_aitoolkit_format(state_dict=state_dict):
                     model = lora_model_from_flux_aitoolkit_state_dict(state_dict=state_dict)
                 else:
                     raise ValueError(f"LoRA model is in unsupported FLUX format: {config.format}")

invokeai/backend/patches/lora_conversions/flux_aitoolkit_lora_conversion_utils.py

@@ -1,5 +1,6 @@
 import json
 from collections import defaultdict
+from dataclasses import dataclass, field
 from typing import Any
 
 import torch
@@ -8,33 +9,44 @@
 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.model_patch_raw import ModelPatchRaw
+from invokeai.backend.util import InvokeAILogger
 
 
-def is_state_dict_likely_in_aitoolkit_format(state_dict: dict[str, Any], metadata: dict[str, Any] = None) -> bool:
+def is_state_dict_likely_in_flux_aitoolkit_format(state_dict: dict[str, Any], metadata: dict[str, Any] = None) -> bool:
     if metadata:
-        software = json.loads(metadata.get("software", "{}"))
+        try:
+            software = json.loads(metadata.get("software", "{}"))
+        except json.JSONDecodeError:
+            return False
         return software.get("name") == "ai-toolkit"
     # metadata got lost somewhere
     return any("diffusion_model" == k.split(".", 1)[0] for k in state_dict.keys())
 
 
-def lora_model_from_flux_aitoolkit_state_dict(state_dict: dict[str, torch.Tensor]) -> ModelPatchRaw:
-    # Group keys by layer.
-    grouped_state_dict: dict[str, dict[str, torch.Tensor]] = defaultdict(dict)
+@dataclass
+class GroupedStateDict:
+    transformer: dict = field(default_factory=dict)
+    # might also grow CLIP and T5 submodels
+
+
+def _group_state_by_submodel(state_dict: dict[str, torch.Tensor]) -> GroupedStateDict:
+    logger = InvokeAILogger.get_logger()
+    grouped = GroupedStateDict()
     for key, value in state_dict.items():
-        layer_name, param_name = key.split(".", 1)
-        grouped_state_dict[layer_name][param_name] = value
+        submodel_name, param_name = key.split(".", 1)
+        match submodel_name:
+            case "diffusion_model":
+                grouped.transformer[param_name] = value
+            case _:
+                logger.warning(f"Unexpected submodel name: {submodel_name}")
+    return grouped
 
-    transformer_grouped_sd: dict[str, dict[str, torch.Tensor]] = {}
 
-    for layer_name, layer_state_dict in grouped_state_dict.items():
-        if layer_name.startswith("diffusion_model"):
-            transformer_grouped_sd[layer_name] = layer_state_dict
-        else:
-            raise ValueError(f"Layer '{layer_name}' does not match the expected pattern for FLUX LoRA weights.")
+def lora_model_from_flux_aitoolkit_state_dict(state_dict: dict[str, torch.Tensor]) -> ModelPatchRaw:
+    grouped = _group_state_by_submodel(state_dict)
 
     layers: dict[str, BaseLayerPatch] = {}
-    for layer_key, layer_state_dict in transformer_grouped_sd.items():
+    for layer_key, layer_state_dict in grouped.transformer.items():
         layers[FLUX_LORA_TRANSFORMER_PREFIX + layer_key] = any_lora_layer_from_state_dict(layer_state_dict)
 
     return ModelPatchRaw(layers=layers)

invokeai/backend/patches/lora_conversions/formats.py

@@ -1,6 +1,6 @@
 from invokeai.backend.model_manager.taxonomy import FluxLoRAFormat
 from invokeai.backend.patches.lora_conversions.flux_aitoolkit_lora_conversion_utils import (
-    is_state_dict_likely_in_aitoolkit_format,
+    is_state_dict_likely_in_flux_aitoolkit_format,
 )
 from invokeai.backend.patches.lora_conversions.flux_control_lora_utils import is_state_dict_likely_flux_control
 from invokeai.backend.patches.lora_conversions.flux_diffusers_lora_conversion_utils import (
@@ -15,7 +15,7 @@
 
 
 def flux_format_from_state_dict(state_dict: dict, metadata: dict | None = None) -> FluxLoRAFormat | None:
-    if is_state_dict_likely_in_aitoolkit_format(state_dict, metadata):
+    if is_state_dict_likely_in_flux_aitoolkit_format(state_dict, metadata):
         return FluxLoRAFormat.AIToolkit
     if is_state_dict_likely_in_flux_kohya_format(state_dict):
         return FluxLoRAFormat.Kohya