fix(LoRA): add ai-toolkit to lora loader

Author: Kevin Turner

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

@@ -20,6 +20,10 @@
     ModelType,
     SubModelType,
 )
+from invokeai.backend.patches.lora_conversions.flux_aitoolkit_lora_conversion_utils import (
+    is_state_dict_likely_in_aitoolkit_format,
+    lora_model_from_flux_aitoolkit_state_dict,
+)
 from invokeai.backend.patches.lora_conversions.flux_control_lora_utils import (
     is_state_dict_likely_flux_control,
     lora_model_from_flux_control_state_dict,
@@ -92,6 +96,8 @@ 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):
+                    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}")
             else:

invokeai/backend/patches/lora_conversions/flux_aitoolkit_lora_conversion_utils.py

@@ -11,15 +11,15 @@
 from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
 
 
-def is_state_dict_likely_in_aitoolkit_format(state_dict: dict[str, Any], metadata: dict[str, Any]) -> bool:
+def is_state_dict_likely_in_aitoolkit_format(state_dict: dict[str, Any], metadata: dict[str, Any] = None) -> bool:
     if metadata:
         software = json.loads(metadata.get("software", "{}"))
         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_aitoolkit_state_dict(state_dict: dict[str, torch.Tensor]) -> ModelPatchRaw:
+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)
     for key, value in state_dict.items():