diff --git a/invokeai/backend/model_manager/configs/lora.py b/invokeai/backend/model_manager/configs/lora.py index 55b4de8683b..d4108c08946 100644 --- a/invokeai/backend/model_manager/configs/lora.py +++ b/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. diff --git a/invokeai/backend/patches/lora_conversions/krea2_lora_conversion_utils.py b/invokeai/backend/patches/lora_conversions/krea2_lora_conversion_utils.py index 458640c906d..566101509de 100644 --- a/invokeai/backend/patches/lora_conversions/krea2_lora_conversion_utils.py +++ b/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 ``.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", } diff --git a/tests/backend/model_manager/configs/test_krea2_lora_config.py b/tests/backend/model_manager/configs/test_krea2_lora_config.py index d9dd0c42e50..d03170d82bb 100644 --- a/tests/backend/model_manager/configs/test_krea2_lora_config.py +++ b/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}) diff --git a/tests/backend/patches/lora_conversions/test_krea2_lora_conversion_utils.py b/tests/backend/patches/lora_conversions/test_krea2_lora_conversion_utils.py index d4ebcc7a431..6e8492c50a1 100644 --- a/tests/backend/patches/lora_conversions/test_krea2_lora_conversion_utils.py +++ b/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)