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

invokeai/backend/model_manager/configs/lora.py

@@ -940,6 +940,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 = (
@@ -983,7 +1010,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):
@@ -1000,9 +1029,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

@@ -248,7 +248,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
 
@@ -339,6 +354,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",
@@ -347,6 +365,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

@@ -273,3 +273,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

@@ -5,12 +5,16 @@
 
 from invokeai.backend.model_manager.load.model_loaders.krea2 import KREA2_TRANSFORMER_CONFIG
 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,
+)
 from tests.backend.patches.lora_conversions.lora_state_dicts.krea2_lora_kohya_format import (
     state_dict_keys as krea2_kohya_state_dict_keys,
 )
@@ -334,25 +338,29 @@ def test_unrecognized_kohya_flattened_keys_are_left_untouched(flat_module: str)
 
 
 @pytest.mark.parametrize(
-    "lycoris_suffixes",
+    ("lycoris_suffixes", "expected_second_layer"),
     [
-        ("lokr_w1", "lokr_w2"),
-        ("hada_w1_a", "hada_w1_b", "hada_w2_a", "hada_w2_b"),
-        ("diff", "diff_b"),
+        # LoKr is a layout this converter understands, so its module un-flattens like any other and the
+        # adapter actually applies. The other algorithms below have no handler, so they stay verbatim.
+        (("lokr_w1", "lokr_w2"), "transformer_blocks.6.attn.to_q"),
+        (("hada_w1_a", "hada_w1_b", "hada_w2_a", "hada_w2_b"), "lora_unet_blocks_6_attn_wq"),
+        (("diff", "diff_b"), "lora_unet_blocks_6_attn_wq"),
         # LyCORIS saves an `alpha` per module, so this is the realistic on-disk shape rather than an
         # edge case — see this repo's own captured fixtures. `.alpha` is a suffix the converter knows,
         # so deciding per key rewrote it while its siblings stayed verbatim, splitting one module into
         # two groups and aborting the load on the orphaned `{'alpha'}`.
-        ("lokr_w1", "lokr_w2", "alpha"),
-        ("hada_w1_a", "hada_w1_b", "hada_w2_a", "hada_w2_b", "alpha"),
+        (("lokr_w1", "lokr_w2", "alpha"), "transformer_blocks.6.attn.to_q"),
+        (("hada_w1_a", "hada_w1_b", "hada_w2_a", "hada_w2_b", "alpha"), "lora_unet_blocks_6_attn_wq"),
         # `dora_scale` is deliberately not combined with a LyCORIS algorithm here: it would orphan the
         # same way, but even grouped correctly `any_lora_layer_from_state_dict` tests `dora_scale`
         # before `lokr_w1`, so a weight-decomposed LoKr dispatches to DoRALayer and dies on a missing
         # `lora_up.weight`. That precedence is shared code, predates this branch, and is not what the
         # per-module gate below is about.
     ],
 )
-def test_kohya_lycoris_algorithm_keys_do_not_abort_the_load(lycoris_suffixes: tuple[str, ...]) -> None:
+def test_kohya_lycoris_algorithm_keys_do_not_abort_the_load(
+    lycoris_suffixes: tuple[str, ...], expected_second_layer: str
+) -> None:
     # LyCORIS supports per-module algorithms, so one kohya file can mix ordinary lora_down/up modules with
     # LoKr/LoHa/full ones. Un-flattening a key whose suffix `_group_by_layer` cannot split back off used to
     # feed it a dotted path, whose blind `rsplit(".", 2)` fallback then cut inside the module name and fused
@@ -368,11 +376,12 @@ def test_kohya_lycoris_algorithm_keys_do_not_abort_the_load(lycoris_suffixes: tu
 
     model = lora_model_from_krea2_state_dict(state_dict)
 
-    # The ordinary module still converts, and the LyCORIS one stays verbatim so it degrades to the per-layer
-    # "Failed to find module" warning at apply time rather than taking the whole adapter down.
+    # Either way the load completes: a supported algorithm converts onto its real module, an unsupported one
+    # stays verbatim and degrades to the per-layer "Failed to find module" warning at apply time, rather than
+    # taking the whole adapter down.
     assert set(model.layers) == {
         f"{KREA2_LORA_TRANSFORMER_PREFIX}transformer_blocks.0.attn.to_v",
-        f"{KREA2_LORA_TRANSFORMER_PREFIX}lora_unet_blocks_6_attn_wq",
+        f"{KREA2_LORA_TRANSFORMER_PREFIX}{expected_second_layer}",
     }
 
 
@@ -447,3 +456,96 @@ 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)
+
+
+def test_kohya_flattened_lokr_converts_onto_its_real_module() -> None:
+    # A LoKr adapter saved in the kohya flattened layout has to clear both hurdles at once: the un-flattening
+    # pass has to reconstruct the dotted module path, and the grouper has to recognise the `lokr_*` suffixes.
+    # Before LoKr was a known suffix the module was left verbatim and the adapter was a silent no-op.
+    state_dict = {
+        "lora_unet_txtfusion_layerwise_blocks_0_attn_wq.lokr_w1": torch.ones(2, 2),
+        "lora_unet_txtfusion_layerwise_blocks_0_attn_wq.lokr_w2": torch.ones(3, 4),
+        "lora_unet_txtfusion_layerwise_blocks_0_attn_wq.alpha": torch.tensor(4.0),
+    }
+
+    model = lora_model_from_krea2_state_dict(state_dict)
+
+    layer = model.layers[f"{KREA2_LORA_TRANSFORMER_PREFIX}text_fusion.layerwise_blocks.0.attn.to_q"]
+    assert isinstance(layer, LoKRLayer)
+    assert layer.get_weight(torch.empty(6, 8)).shape == (6, 8)