feat(model-manager): support LyCORIS LoKr LoRAs for Krea-2

invokeai/backend/model_manager/configs/lora.py

@@ -1018,6 +1018,33 @@ def _has_complete_lora_pair(state_dict: dict[str | int, Any], prefixes: tuple[st
     return False
 
 
+# LyCORIS LoKr layers carry their Kronecker factors instead of a lora_A/B (or lora_down/up) pair: either the
+# full `lokr_w1`/`lokr_w2`, or the further-factored `lokr_w1_a`/`lokr_w1_b` / `lokr_w2_a`/`lokr_w2_b` (+ the
+# optional `lokr_t2` tucker core). Each such layer is self-contained, so there is no "orphaned half" notion to
+# check — presence of any factor is enough to call the layer complete.
+_LOKR_WEIGHT_SUFFIXES = (
+    ".lokr_w1",
+    ".lokr_w2",
+    ".lokr_w1_a",
+    ".lokr_w1_b",
+    ".lokr_w2_a",
+    ".lokr_w2_b",
+    ".lokr_t2",
+)
+
+
+def _has_lokr_layer(state_dict: dict[str | int, Any], prefixes: tuple[str, ...] | None = None) -> bool:
+    """True if the state dict contains at least one LoKr layer, optionally restricted to `prefixes`."""
+    for key in state_dict:
+        if not isinstance(key, str):
+            continue
+        if prefixes is not None and not key.startswith(prefixes):
+            continue
+        if key.endswith(_LOKR_WEIGHT_SUFFIXES):
+            return True
+    return False
+
+
 # Layouts the converter understands for an explicit Krea-2 override (a transformer-only or text-encoder-only
 # LoRA that lacks the auto-detection text_fusion/time_mod_proj keys still installs under an explicit base).
 _KREA2_SUPPORTED_LORA_PREFIXES = (
@@ -1051,7 +1078,9 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -
 
         state_dict = mod.load_state_dict()
         explicit_krea2_override = override_fields.get("base") is BaseModelType.Krea2
-        has_supported_explicit_pair = _has_complete_lora_pair(state_dict, _KREA2_SUPPORTED_LORA_PREFIXES)
+        has_supported_explicit_pair = _has_complete_lora_pair(
+            state_dict, _KREA2_SUPPORTED_LORA_PREFIXES
+        ) or _has_lokr_layer(state_dict, _KREA2_SUPPORTED_LORA_PREFIXES)
         # Reject an orphaned half *anywhere* in the state dict (e.g. a dangling text_fusion half not under
         # the approved prefixes) — it would install here but fail during LoRA conversion at generation time.
         if explicit_krea2_override and has_supported_explicit_pair and _lora_weight_keys_are_all_paired(state_dict):
@@ -1068,9 +1097,11 @@ def _validate_looks_like_lora(cls, mod: ModelOnDisk) -> None:
         state_dict = mod.load_state_dict()
         # Require a *complete* lora_A/B (or lora_down/up) pair, not merely any lora/dora suffix: a file with
         # only ``dora_scale`` and no A/B weights would pass a suffix check but fail later on missing weights.
-        if not (_has_krea2_lora_keys(state_dict) and _has_complete_lora_pair(state_dict)):
+        if not (
+            _has_krea2_lora_keys(state_dict) and (_has_complete_lora_pair(state_dict) or _has_lokr_layer(state_dict))
+        ):
             raise NotAMatchError(
-                "model does not match Krea-2 LoRA heuristics (no complete lora_A/B or lora_down/up pair)"
+                "model does not match Krea-2 LoRA heuristics (no complete lora_A/B, lora_down/up or LoKr layer)"
             )
         # Reject a file with an orphaned LoRA half (a valid layer plus a dangling lora_A/B/down/up); it
         # would install here but fail later during LoRA conversion.

invokeai/backend/patches/lora_conversions/krea2_lora_conversion_utils.py

@@ -113,7 +113,22 @@ def is_state_dict_likely_krea2_lora(state_dict: dict[str | int, torch.Tensor]) -
     str_keys = [k for k in state_dict.keys() if isinstance(k, str)]
     has_krea2_module = any(any(sig in k for sig in KREA2_TRANSFORMER_SIGNATURE_KEYS) for k in str_keys)
     has_lora_suffix = any(
-        k.endswith((".lora_A.weight", ".lora_B.weight", ".lora_down.weight", ".lora_up.weight")) for k in str_keys
+        k.endswith(
+            (
+                ".lora_A.weight",
+                ".lora_B.weight",
+                ".lora_down.weight",
+                ".lora_up.weight",
+                # LyCORIS LoKr (e.g. ai-toolkit Krea-2 adapters) stores Kronecker factors, not an A/B pair.
+                ".lokr_w1",
+                ".lokr_w2",
+                ".lokr_w1_a",
+                ".lokr_w1_b",
+                ".lokr_w2_a",
+                ".lokr_w2_b",
+            )
+        )
+        for k in str_keys
     )
     return has_krea2_module and has_lora_suffix
 
@@ -202,6 +217,9 @@ def _get_lora_layer_values(
 # magnitude is published as ``<layer>.lora_magnitude_vector.weight``; it is the same thing InvokeAI stores as
 # ``dora_scale``, so mapping it here lets a standard Diffusers DoRA adapter (A/B + magnitude) load as a
 # DoRALayer instead of being split into a bogus, unrecognized layer.
+#
+# LyCORIS LoKr factors are passed through untouched: `any_lora_layer_from_state_dict` routes a values dict
+# containing `lokr_w1` / `lokr_w1_a` to LoKRLayer, so they only need to survive _group_by_layer intact.
 _SUFFIX_TO_VALUE_KEY = {
     ".lora_A.weight": "lora_A.weight",
     ".lora_B.weight": "lora_B.weight",
@@ -210,6 +228,13 @@ def _get_lora_layer_values(
     ".dora_scale": "dora_scale",
     ".lora_magnitude_vector.weight": "dora_scale",
     ".alpha": "alpha",
+    ".lokr_w1": "lokr_w1",
+    ".lokr_w2": "lokr_w2",
+    ".lokr_w1_a": "lokr_w1_a",
+    ".lokr_w1_b": "lokr_w1_b",
+    ".lokr_w2_a": "lokr_w2_a",
+    ".lokr_w2_b": "lokr_w2_b",
+    ".lokr_t2": "lokr_t2",
 }
 
 

tests/backend/model_manager/configs/test_krea2_lora_config.py

@@ -237,3 +237,45 @@ def test_explicit_krea2_override_accepts_single_module_native_lora(_raise_if_not
     config = LoRA_LyCORIS_Krea2_Config.from_model_on_disk(mod, {**_REQUIRED_FIELDS, "base": BaseModelType.Krea2})
 
     assert config.base is BaseModelType.Krea2
+
+
+def _native_lokr_lora() -> MagicMock:
+    # LyCORIS LoKr adapter targeting the Krea-2 text-fusion stage (the layout ai-toolkit emits). It carries
+    # Kronecker factors instead of a lora_A/lora_B pair.
+    mod = MagicMock()
+    mod.load_state_dict.return_value = {
+        "diffusion_model.txtfusion.layerwise_blocks.0.attn.wq.lokr_w1": object(),
+        "diffusion_model.txtfusion.layerwise_blocks.0.attn.wq.lokr_w2": object(),
+        "diffusion_model.txtfusion.layerwise_blocks.0.attn.wq.alpha": object(),
+    }
+    return mod
+
+
+@patch("invokeai.backend.model_manager.configs.lora.raise_if_not_file")
+def test_automatic_probe_accepts_lokr_lora(_raise_if_not_file) -> None:
+    config = LoRA_LyCORIS_Krea2_Config.from_model_on_disk(_native_lokr_lora(), {**_REQUIRED_FIELDS})
+
+    assert config.base is BaseModelType.Krea2
+
+
+@patch("invokeai.backend.model_manager.configs.lora.raise_if_not_file")
+def test_explicit_krea2_override_accepts_lokr_lora(_raise_if_not_file) -> None:
+    config = LoRA_LyCORIS_Krea2_Config.from_model_on_disk(
+        _native_lokr_lora(), {**_REQUIRED_FIELDS, "base": BaseModelType.Krea2}
+    )
+
+    assert config.base is BaseModelType.Krea2
+
+
+@patch("invokeai.backend.model_manager.configs.lora.raise_if_not_file")
+def test_automatic_probe_rejects_lokr_without_krea2_modules(_raise_if_not_file) -> None:
+    # A LoKr that does not touch the Krea-2 signature modules belongs to another base and must not be
+    # claimed here just because it is a LoKr.
+    mod = MagicMock()
+    mod.load_state_dict.return_value = {
+        "transformer.transformer_blocks.0.attn.to_q.lokr_w1": object(),
+        "transformer.transformer_blocks.0.attn.to_q.lokr_w2": object(),
+    }
+
+    with pytest.raises(NotAMatchError):
+        LoRA_LyCORIS_Krea2_Config.from_model_on_disk(mod, {**_REQUIRED_FIELDS})

tests/backend/patches/lora_conversions/test_krea2_lora_conversion_utils.py

@@ -2,12 +2,16 @@
 import torch
 
 from invokeai.backend.patches.layers.dora_layer import DoRALayer
+from invokeai.backend.patches.layers.lokr_layer import LoKRLayer
 from invokeai.backend.patches.layers.lora_layer import LoRALayer
 from invokeai.backend.patches.lora_conversions.krea2_lora_constants import (
     KREA2_LORA_QWEN3VL_PREFIX,
     KREA2_LORA_TRANSFORMER_PREFIX,
 )
-from invokeai.backend.patches.lora_conversions.krea2_lora_conversion_utils import lora_model_from_krea2_state_dict
+from invokeai.backend.patches.lora_conversions.krea2_lora_conversion_utils import (
+    is_state_dict_likely_krea2_lora,
+    lora_model_from_krea2_state_dict,
+)
 
 
 def test_peft_layer_preserves_explicit_alpha() -> None:
@@ -250,3 +254,79 @@ def test_native_krea2_top_level_linear_keys_are_remapped() -> None:
         f"{KREA2_LORA_TRANSFORMER_PREFIX}{diffusers_module}" for diffusers_module in native_to_diffusers.values()
     }
     assert expected_keys < set(model.layers)
+
+
+def test_lokr_layer_produces_lokr_layer() -> None:
+    # LyCORIS LoKr adapters (e.g. those produced by ai-toolkit for Krea-2) carry Kronecker factors instead of
+    # a lora_A/lora_B pair. They must survive _group_by_layer intact so any_lora_layer_from_state_dict can
+    # route them to LoKRLayer.
+    state_dict = {
+        "transformer.text_fusion.0.attn.to_q.lokr_w1": torch.ones(2, 2),
+        "transformer.text_fusion.0.attn.to_q.lokr_w2": torch.ones(3, 4),
+        "transformer.text_fusion.0.attn.to_q.alpha": torch.tensor(1.0),
+    }
+
+    model = lora_model_from_krea2_state_dict(state_dict)
+
+    layer = model.layers[f"{KREA2_LORA_TRANSFORMER_PREFIX}text_fusion.0.attn.to_q"]
+    assert isinstance(layer, LoKRLayer)
+    assert layer._alpha == 1.0
+    # The reconstructed weight is the Kronecker product of the two factors.
+    assert layer.get_weight(torch.empty(6, 8)).shape == (6, 8)
+
+
+def test_factored_lokr_layer_produces_lokr_layer() -> None:
+    # LoKr may factor either Kronecker operand further into an `_a`/`_b` pair. Both spellings must be grouped
+    # onto the same layer.
+    state_dict = {
+        "transformer.text_fusion.0.attn.to_q.lokr_w1_a": torch.ones(2, 1),
+        "transformer.text_fusion.0.attn.to_q.lokr_w1_b": torch.ones(1, 2),
+        "transformer.text_fusion.0.attn.to_q.lokr_w2_a": torch.ones(3, 1),
+        "transformer.text_fusion.0.attn.to_q.lokr_w2_b": torch.ones(1, 4),
+    }
+
+    model = lora_model_from_krea2_state_dict(state_dict)
+
+    layer = model.layers[f"{KREA2_LORA_TRANSFORMER_PREFIX}text_fusion.0.attn.to_q"]
+    assert isinstance(layer, LoKRLayer)
+    assert layer.w1_a is not None and layer.w1_b is not None
+    assert layer.w2_a is not None and layer.w2_b is not None
+
+
+def test_native_lokr_keys_are_renamed_to_diffusers_layout() -> None:
+    # Native (ComfyUI / ai-toolkit) LoKr keys must go through the same native->diffusers renaming as LoRA
+    # keys: txtfusion -> text_fusion, attn.wq -> attn.to_q, mlp.down -> ff.down.
+    state_dict = {
+        "diffusion_model.txtfusion.layerwise_blocks.0.attn.wq.lokr_w1": torch.ones(2, 2),
+        "diffusion_model.txtfusion.layerwise_blocks.0.attn.wq.lokr_w2": torch.ones(3, 4),
+        "diffusion_model.txtfusion.refiner_blocks.1.mlp.down.lokr_w1": torch.ones(2, 2),
+        "diffusion_model.txtfusion.refiner_blocks.1.mlp.down.lokr_w2": torch.ones(3, 4),
+    }
+
+    model = lora_model_from_krea2_state_dict(state_dict)
+
+    assert set(model.layers) == {
+        f"{KREA2_LORA_TRANSFORMER_PREFIX}text_fusion.layerwise_blocks.0.attn.to_q",
+        f"{KREA2_LORA_TRANSFORMER_PREFIX}text_fusion.refiner_blocks.1.ff.down",
+    }
+    assert all(isinstance(layer, LoKRLayer) for layer in model.layers.values())
+
+
+def test_is_state_dict_likely_krea2_lora_accepts_lokr() -> None:
+    state_dict = {
+        "diffusion_model.txtfusion.layerwise_blocks.0.attn.wq.lokr_w1": torch.ones(2, 2),
+        "diffusion_model.txtfusion.layerwise_blocks.0.attn.wq.lokr_w2": torch.ones(3, 4),
+    }
+
+    assert is_state_dict_likely_krea2_lora(state_dict)
+
+
+def test_is_state_dict_likely_krea2_lora_rejects_lokr_without_krea2_modules() -> None:
+    # The Krea-2 signature modules are still required: a LoKr targeting only generic transformer blocks
+    # belongs to another base (e.g. Qwen-Image) and must not be claimed here.
+    state_dict = {
+        "transformer.transformer_blocks.0.attn.to_q.lokr_w1": torch.ones(2, 2),
+        "transformer.transformer_blocks.0.attn.to_q.lokr_w2": torch.ones(3, 4),
+    }
+
+    assert not is_state_dict_likely_krea2_lora(state_dict)