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

invokeai/backend/model_manager/configs/lora.py

@@ -48,7 +48,7 @@
     has_cosmos_dit_peft_keys_strict,
 )
 from invokeai.backend.patches.lora_conversions.flux_control_lora_utils import is_state_dict_likely_flux_control
-from invokeai.backend.patches.lora_conversions.krea2_lora_constants import is_kohya_krea2_lora_key
+from invokeai.backend.patches.lora_conversions.krea2_lora_constants import LOKR_WEIGHT_KEYS, is_kohya_krea2_lora_key
 from invokeai.backend.patches.lora_conversions.ltx2_lora_constants import (
     has_ltx2_lora_keys,
     has_non_ltx2_architecture_keys,
@@ -996,6 +996,56 @@ def _has_complete_lora_pair(state_dict: dict[str | int, Any], key_filter: Callab
     return False
 
 
+# LyCORIS LoKr layers carry Kronecker operands instead of a lora_A/B (or lora_down/up) pair. `LoKRLayer` needs
+# both operands, each exactly once: whole (`lokr_w1` / `lokr_w2`) or factored into a complete `_a`/`_b` pair.
+# A module missing an operand, holding half a factored pair, or spelling one operand both ways installs fine
+# but fails when the layer is built at generation time — the LoKr counterpart of an orphaned lora_A/B half —
+# so completeness is checked per module, like the A/B pairing.
+_LOKR_WEIGHT_SUFFIXES = tuple(f".{key}" for key in LOKR_WEIGHT_KEYS)
+
+
+def _lokr_suffixes_by_module(
+    state_dict: dict[str | int, Any], key_filter: Callable[[str], bool] | None = None
+) -> dict[str, set[str]]:
+    """Group the LoKr weight keys (optionally restricted to the keys `key_filter` accepts) by module path."""
+    suffixes_by_module: dict[str, set[str]] = {}
+    for key in state_dict:
+        if not isinstance(key, str) or not key.endswith(_LOKR_WEIGHT_SUFFIXES):
+            continue
+        if key_filter is not None and not key_filter(key):
+            continue
+        module_path, _, suffix = key.rpartition(".")
+        suffixes_by_module.setdefault(module_path, set()).add(suffix)
+    return suffixes_by_module
+
+
+def _lokr_operand_is_complete(suffixes: set[str], operand: str) -> bool:
+    """True if a LoKr operand (`lokr_w1` or `lokr_w2`) is present exactly once: whole, or as a full `_a`/`_b` pair."""
+    whole = operand in suffixes
+    factored = suffixes & {f"{operand}_a", f"{operand}_b"}
+    return (whole and not factored) or (not whole and len(factored) == 2)
+
+
+def _lokr_module_is_complete(suffixes: set[str]) -> bool:
+    return _lokr_operand_is_complete(suffixes, "lokr_w1") and _lokr_operand_is_complete(suffixes, "lokr_w2")
+
+
+def _has_complete_lokr_layer(state_dict: dict[str | int, Any], key_filter: Callable[[str], bool] | None = None) -> bool:
+    """True if at least one complete LoKr module exists, optionally restricted to the keys `key_filter` accepts.
+
+    Like :func:`_has_complete_lora_pair`, this does not reject incomplete modules elsewhere — callers pair it
+    with :func:`_lokr_layers_are_all_complete`.
+    """
+    return any(
+        _lokr_module_is_complete(suffixes) for suffixes in _lokr_suffixes_by_module(state_dict, key_filter).values()
+    )
+
+
+def _lokr_layers_are_all_complete(state_dict: dict[str | int, Any]) -> bool:
+    """True if *every* LoKr module in the state dict is complete. Returns True when there are none."""
+    return all(_lokr_module_is_complete(suffixes) for suffixes in _lokr_suffixes_by_module(state_dict).values())
+
+
 # Dotted 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). The kohya/LyCORIS layout is deliberately absent - see `_key_is_supported_krea2_layout`.
@@ -1044,10 +1094,18 @@ 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, _key_is_supported_krea2_layout)
-        # 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):
+        has_supported_explicit_layer = _has_complete_lora_pair(
+            state_dict, _key_is_supported_krea2_layout
+        ) or _has_complete_lokr_layer(state_dict, _key_is_supported_krea2_layout)
+        # Reject an orphaned half or incomplete LoKr module *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_layer
+            and _lora_weight_keys_are_all_paired(state_dict)
+            and _lokr_layers_are_all_complete(state_dict)
+        ):
             return cls(**override_fields)
 
         cls._validate_looks_like_lora(mod)
@@ -1057,18 +1115,25 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -
     @classmethod
     def _validate_looks_like_lora(cls, mod: ModelOnDisk) -> None:
         """Krea-2 LoRAs have keys like transformer.text_fusion.* / transformer.transformer_blocks.* with
-        a lora_A/lora_B (or lora_down/lora_up) suffix. The text-fusion stage is unique to Krea-2."""
+        a lora_A/lora_B (or lora_down/lora_up) suffix, or LyCORIS LoKr ``lokr_*`` factors. The text-fusion
+        stage is unique to Krea-2."""
         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)):
+        # Require a *complete* lora_A/B (or lora_down/up) pair or LoKr module, not merely any lora/dora/lokr
+        # 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) or _has_complete_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.
+        # Reject a file with an orphaned LoRA half (a valid layer plus a dangling lora_A/B/down/up) or an
+        # incomplete LoKr module; it would install here but fail later during LoRA conversion.
         if not _lora_weight_keys_are_all_paired(state_dict):
             raise NotAMatchError("Krea-2 LoRA has an incomplete lora_A/B (or lora_down/up) weight pair")
+        if not _lokr_layers_are_all_complete(state_dict):
+            raise NotAMatchError("Krea-2 LoRA has an incomplete LoKr module (lokr_w1 / lokr_w2 operands)")
 
     @classmethod
     def _get_base_or_raise(cls, mod: ModelOnDisk) -> BaseModelType:

invokeai/backend/patches/lora_conversions/krea2_lora_constants.py

@@ -20,6 +20,11 @@
 # Prefix for Krea-2 Qwen3-VL text encoder LoRA layers.
 KREA2_LORA_QWEN3VL_PREFIX = "lora_qwen3vl-"
 
+# LyCORIS LoKr weight keys: the Kronecker operands ``lokr_w1`` / ``lokr_w2``, each optionally factored further into
+# an ``_a``/``_b`` pair, plus the optional Tucker core ``lokr_t2``. Identification (``model_manager.configs.lora``)
+# and the converter's suffix map both derive from this one list so they cannot drift apart.
+LOKR_WEIGHT_KEYS = ("lokr_w1", "lokr_w1_a", "lokr_w1_b", "lokr_w2", "lokr_w2_a", "lokr_w2_b", "lokr_t2")
+
 
 # --- Kohya / LyCORIS (flattened) -> native key mapping ---------------------------------------------------------
 # sd-scripts and LyCORIS flatten the module path (``path.replace(".", "_")``) and prefix it with

invokeai/backend/patches/lora_conversions/krea2_lora_conversion_utils.py

@@ -20,10 +20,14 @@
 from invokeai.backend.patches.lora_conversions.krea2_lora_constants import (
     KREA2_LORA_QWEN3VL_PREFIX,
     KREA2_LORA_TRANSFORMER_PREFIX,
+    LOKR_WEIGHT_KEYS,
     split_kohya_krea2_key,
     unflatten_kohya_krea2_module_path,
 )
 from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
+from invokeai.backend.util.logging import InvokeAILogger
+
+logger = InvokeAILogger.get_logger(__name__)
 
 # Module-name fragments unique to the Krea-2 transformer (text-fusion stage + timestep modulation proj).
 # ``text_fusion`` is the diffusers name; ``txtfusion`` is the native (ComfyUI) name for the same stage.
@@ -117,7 +121,8 @@ def _maybe_convert_kohya_krea2_state_dict(
 ) -> Dict[str, torch.Tensor]:
     """Rewrite kohya/LyCORIS flattened Krea-2 keys to the dotted native layout, leaving all others untouched."""
     # Whether a module can be converted is a property of the *module*, not of each key on its own. A LyCORIS
-    # module (``lokr_w1``, ``hada_w1_a``, ``diff``) has to stay verbatim — see below — but LyCORIS also saves
+    # module this converter has no handler for (``hada_w1_a``, ``diff``) has to stay verbatim — see below; LoKr
+    # is the exception, its ``lokr_*`` suffixes are known and convert like any other — but LyCORIS also saves
     # a sibling ``.alpha`` (and ``.dora_scale`` for the weight-decomposed variants), whose suffix this
     # converter *does* recognize. Deciding per key rewrites that sibling alone and splits one module across
     # two layer groups: the rewritten ``.alpha`` ends up in a group by itself, and ``_get_lora_layer_values``
@@ -146,9 +151,9 @@ def _maybe_convert_kohya_krea2_state_dict(
             # Only rewrite when ``_group_by_layer`` can split the suffix back off. Un-flattening introduces
             # dots into the module path, and the grouper's fallback for an unknown suffix is a blind
             # ``rsplit(".", 2)`` — on a dotted path that cuts *inside the module name*, fusing two modules
-            # into one bogus layer that aborts the whole load. LyCORIS suffixes such as ``.lokr_w1`` or
-            # ``.hada_w1_a`` hit exactly that. Flattened, they have no interior dot and group harmlessly,
-            # so leaving them verbatim keeps them at the pre-existing warn-and-skip behaviour.
+            # into one bogus layer that aborts the whole load. LyCORIS suffixes without a handler here, such as
+            # ``.hada_w1_a`` or ``.diff``, hit exactly that. Flattened, they have no interior dot and group
+            # harmlessly, so leaving them verbatim keeps them at the pre-existing warn-and-skip behaviour.
             if module_path is not None and flat_path in fully_convertible_flat_paths:
                 converted_key = f"{module_path}{dot}{weight_suffix}"
         if converted_key in converted_state_dict:
@@ -237,7 +242,8 @@ def lora_model_from_krea2_state_dict(state_dict: Dict[str, torch.Tensor], alpha:
 def _get_lora_layer_values(
     layer_key: str, layer_dict: dict[str, torch.Tensor], alpha: float | None
 ) -> dict[str, torch.Tensor]:
-    """Convert PEFT (lora_A/lora_B) layer values to internal (lora_down/lora_up) format."""
+    """Convert PEFT (lora_A/lora_B) layer values to internal (lora_down/lora_up) format, and drop a DoRA
+    magnitude from a LoKr layer (see below)."""
     if "lora_A.weight" in layer_dict:
         if "lora_B.weight" not in layer_dict:
             raise ValueError(
@@ -256,6 +262,18 @@ def _get_lora_layer_values(
         if alpha is not None:
             values["alpha"] = torch.tensor(alpha)
         return values
+    # A weight-decomposed LoKr (LyCORIS `dora_wd`) also carries a DoRA magnitude. `any_lora_layer_from_state_dict`
+    # tests for that before the LoKr factors and would build a DoRALayer, which fails on the missing
+    # lora_up/lora_down. There is no DoRA+LoKr layer, so apply the plain LoKr delta instead, as the Anima
+    # converter does.
+    if "lokr_w1" in layer_dict or "lokr_w1_a" in layer_dict:
+        magnitude_keys = {"dora_scale", "dora_magnitude"} & layer_dict.keys()
+        if magnitude_keys:
+            logger.warning(
+                f"Krea-2 LoRA layer '{layer_key}' combines LoKr with a DoRA magnitude, which is not supported; "
+                "applying it as plain LoKr."
+            )
+            return {key: value for key, value in layer_dict.items() if key not in magnitude_keys}
     return layer_dict
 
 
@@ -267,6 +285,9 @@ def _get_lora_layer_values(
 #     dim -> value key ``dora_magnitude``
 # Mapping them here lets a 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",
@@ -276,6 +297,7 @@ def _get_lora_layer_values(
     ".lora_magnitude_vector.weight": "dora_magnitude",
     ".magnitude": "dora_magnitude",
     ".alpha": "alpha",
+    **{f".{key}": key for key in LOKR_WEIGHT_KEYS},
 }
 
 

tests/backend/model_manager/configs/test_krea2_lora_config.py

@@ -371,3 +371,118 @@ 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
+
+
+_LOKR_MODULE = "diffusion_model.txtfusion.layerwise_blocks.0.attn.wq"
+
+
+def _lokr_mod(keys: dict[str, object]) -> MagicMock:
+    mod = MagicMock()
+    mod.load_state_dict.return_value = keys
+    return mod
+
+
+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.
+    return _lokr_mod({f"{_LOKR_MODULE}.{suffix}": object() for suffix in ("lokr_w1", "lokr_w2", "alpha")})
+
+
+@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
+
+
+@pytest.mark.parametrize(
+    "suffixes",
+    [
+        ("lokr_w1", "lokr_w2_a", "lokr_w2_b", "alpha"),
+        ("lokr_w1_a", "lokr_w1_b", "lokr_w2", "alpha"),
+        ("lokr_w1_a", "lokr_w1_b", "lokr_w2_a", "lokr_w2_b", "lokr_t2", "alpha"),
+        # The DoRA magnitude of a weight-decomposed LoKr is dropped by the converter, so the file installs.
+        ("lokr_w1", "lokr_w2", "dora_scale", "alpha"),
+    ],
+)
+@patch("invokeai.backend.model_manager.configs.lora.raise_if_not_file")
+def test_automatic_probe_accepts_every_complete_lokr_spelling(_raise_if_not_file, suffixes: tuple[str, ...]) -> None:
+    mod = _lokr_mod({f"{_LOKR_MODULE}.{suffix}": object() for suffix in suffixes})
+
+    config = LoRA_LyCORIS_Krea2_Config.from_model_on_disk(mod, {**_REQUIRED_FIELDS})
+
+    assert config.base is BaseModelType.Krea2
+
+
+@pytest.mark.parametrize(
+    "suffixes",
+    [
+        # Each of these used to install and then fail building LoKRLayer at generation time.
+        ("lokr_w1", "alpha"),
+        ("lokr_w2", "alpha"),
+        ("lokr_w1", "lokr_w2_a", "alpha"),
+        ("lokr_w1_a", "lokr_w2", "alpha"),
+        # Both spellings of one operand: LoKRLayer requires exactly one.
+        ("lokr_w1", "lokr_w1_a", "lokr_w1_b", "lokr_w2", "alpha"),
+        ("lokr_t2", "alpha"),
+    ],
+)
+@pytest.mark.parametrize("override_base", [False, True])
+@patch("invokeai.backend.model_manager.configs.lora.raise_if_not_file")
+def test_incomplete_lokr_module_is_rejected(_raise_if_not_file, suffixes: tuple[str, ...], override_base: bool) -> None:
+    # A complete module elsewhere in the file does not excuse the incomplete one: like an orphaned lora_A/B
+    # half, it would take the whole adapter down at generation time.
+    keys: dict[str, object] = {f"{_LOKR_MODULE}.{suffix}": object() for suffix in suffixes}
+    keys.update({f"diffusion_model.blocks.0.attn.wq.{suffix}": object() for suffix in ("lokr_w1", "lokr_w2")})
+    override = {"base": BaseModelType.Krea2} if override_base else {}
+
+    with pytest.raises(NotAMatchError):
+        LoRA_LyCORIS_Krea2_Config.from_model_on_disk(_lokr_mod(keys), {**_REQUIRED_FIELDS, **override})
+
+
+@pytest.mark.parametrize(
+    "module",
+    [
+        # Dotted layout without the Krea-2 signature modules (no text fusion, no gated attention pair).
+        "transformer.transformer_blocks.0.attn.to_q",
+        # kohya-flattened Krea-2 module; the automatic probe only reads dotted names.
+        "lora_unet_blocks_0_attn_wq",
+    ],
+)
+@patch("invokeai.backend.model_manager.configs.lora.raise_if_not_file")
+def test_explicit_krea2_override_accepts_lokr_the_probe_cannot_identify(_raise_if_not_file, module: str) -> None:
+    keys = {f"{module}.{suffix}": object() for suffix in ("lokr_w1", "lokr_w2", "alpha")}
+
+    with pytest.raises(NotAMatchError):
+        LoRA_LyCORIS_Krea2_Config.from_model_on_disk(_lokr_mod(keys), {**_REQUIRED_FIELDS})
+
+    config = LoRA_LyCORIS_Krea2_Config.from_model_on_disk(
+        _lokr_mod(keys), {**_REQUIRED_FIELDS, "base": BaseModelType.Krea2}
+    )
+
+    assert config.base is BaseModelType.Krea2
+
+
+@patch("invokeai.backend.model_manager.configs.lora.raise_if_not_file")
+def test_explicit_krea2_override_rejects_wan_kohya_lokr(_raise_if_not_file) -> None:
+    # Wan writes the same `lora_unet_blocks_<idx>_` spelling, but `self_attn_q` is no Krea-2 module, so forcing
+    # the base must not let it install here and silently no-op at generation time.
+    keys = {f"lora_unet_blocks_0_self_attn_q.{suffix}": object() for suffix in ("lokr_w1", "lokr_w2", "alpha")}
+
+    with pytest.raises(NotAMatchError):
+        LoRA_LyCORIS_Krea2_Config.from_model_on_disk(_lokr_mod(keys), {**_REQUIRED_FIELDS, "base": 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 = _lokr_mod(
+        {
+            "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})