From 7b4f0f4953fcbdd6fde22e5f14cfc437d67306f8 Mon Sep 17 00:00:00 2001 From: Balrog994 Date: Fri, 31 Jul 2026 21:51:48 +0200 Subject: [PATCH 1/2] feat(model-manager): support LyCORIS LoKr LoRAs for Krea-2 Krea-2 LoRAs in LyCORIS LoKr format were rejected at install time and could not be used. LoKr is the default output format of ai-toolkit, so this affects a fair number of community adapters. Two gaps in the Krea-2 path: - `LoRA_LyCORIS_Krea2_Config` required a complete lora_A/B or lora_down/up pair. LoKr carries Kronecker factors instead, so the probe rejected it even though the Krea-2 signature modules matched. The explicit-base override went through the same helper, so forcing the base did not help either. - `_SUFFIX_TO_VALUE_KEY` did not map the lokr_* suffixes, so `_group_by_layer` would have dropped the weights and left each layer holding only its alpha. The rest of the stack already handles LoKr: `LoKRLayer` exists and `any_lora_layer_from_state_dict` routes to it on `lokr_w1` / `lokr_w1_a`. Only the Krea-2 wiring was missing, matching its (WIP) status in #9304. The Krea-2 signature-module requirement is unchanged, so a LoKr targeting only generic transformer blocks is still not claimed by this base. --- .../backend/model_manager/configs/lora.py | 37 ++++++++- .../krea2_lora_conversion_utils.py | 27 +++++- .../configs/test_krea2_lora_config.py | 42 ++++++++++ .../test_krea2_lora_conversion_utils.py | 82 ++++++++++++++++++- 4 files changed, 183 insertions(+), 5 deletions(-) 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) From 6de2785e871c164a3c50192c37154edf7eb057ec Mon Sep 17 00:00:00 2001 From: Balrog994 Date: Mon, 12 Oct 2026 22:02:35 +0200 Subject: [PATCH 2/2] fix(krea2-lora): require complete LoKr modules, pin the factor mapping in tests Addresses review 9424: - Identification checked only that some `lokr_*` key existed, so a module missing an operand, holding half a factored `_a`/`_b` pair, or spelling one operand both ways installed and then failed building LoKRLayer at generation time. Completeness is now checked per module, mirroring LoKRLayer's own invariants, on both the automatic and the explicit-override path. - A weight-decomposed LoKr (`dora_scale` / `magnitude` beside the factors) dispatched to DoRALayer and died on the missing lora_up/lora_down. The converter now drops the magnitude and applies plain LoKr with a warning, as the Anima converter does. - The LoKr suffix list lives once, in krea2_lora_constants, and both the identification helpers and `_SUFFIX_TO_VALUE_KEY` derive from it. - `is_state_dict_likely_krea2_lora` has no callers; its LoKr change and the two tests for it are reverted. Stale comments in the kohya pass and in `_validate_looks_like_lora` are updated. Tests: the override tests now use LoKr files the automatic probe cannot identify (dotted without the Krea-2 signature, and kohya-flattened), plus a Wan kohya LoKr that must still be refused. The conversion tests use random factors and compare against torch.kron exactly, and check scale() == alpha / rank for the factored case, so swapping w1/w2 or _a/_b in the suffix map now fails them. Co-Authored-By: Claude Opus 5.5 --- .../backend/model_manager/configs/lora.py | 97 ++++++++++----- .../lora_conversions/krea2_lora_constants.py | 5 + .../krea2_lora_conversion_utils.py | 53 ++++----- .../configs/test_krea2_lora_config.py | 101 +++++++++++++--- .../test_krea2_lora_conversion_utils.py | 111 ++++++++++-------- 5 files changed, 241 insertions(+), 126 deletions(-) diff --git a/invokeai/backend/model_manager/configs/lora.py b/invokeai/backend/model_manager/configs/lora.py index abcd05a54f0..e46cc034a14 100644 --- a/invokeai/backend/model_manager/configs/lora.py +++ b/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, @@ -956,32 +956,54 @@ def _has_complete_lora_pair(state_dict: dict[str | int, Any], key_filter: Callab 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", -) +# 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 _has_lokr_layer(state_dict: dict[str | int, Any], key_filter: Callable[[str], bool] | None = None) -> bool: - """True if the state dict contains at least one LoKr layer, optionally restricted to the keys - `key_filter` accepts.""" +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): + 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 - if key.endswith(_LOKR_WEIGHT_SUFFIXES): - return True - return False + 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 @@ -1032,12 +1054,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( + has_supported_explicit_layer = _has_complete_lora_pair( state_dict, _key_is_supported_krea2_layout - ) or _has_lokr_layer(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): + ) 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) @@ -1047,20 +1075,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. + # 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_lokr_layer(state_dict)) + _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, 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: diff --git a/invokeai/backend/patches/lora_conversions/krea2_lora_constants.py b/invokeai/backend/patches/lora_conversions/krea2_lora_constants.py index 433a4d1242d..50dc1632d0b 100644 --- a/invokeai/backend/patches/lora_conversions/krea2_lora_constants.py +++ b/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 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 0267bcdbfce..f4341c241e8 100644 --- a/invokeai/backend/patches/lora_conversions/krea2_lora_conversion_utils.py +++ b/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: @@ -171,22 +176,7 @@ 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", - # 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 + k.endswith((".lora_A.weight", ".lora_B.weight", ".lora_down.weight", ".lora_up.weight")) for k in str_keys ) return has_krea2_module and has_lora_suffix @@ -252,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( @@ -271,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 @@ -294,13 +297,7 @@ def _get_lora_layer_values( ".lora_magnitude_vector.weight": "dora_magnitude", ".magnitude": "dora_magnitude", ".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", + **{f".{key}": key for key in LOKR_WEIGHT_KEYS}, } 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 9972165b231..fe7fa052312 100644 --- a/tests/backend/model_manager/configs/test_krea2_lora_config.py +++ b/tests/backend/model_manager/configs/test_krea2_lora_config.py @@ -373,16 +373,19 @@ def test_explicit_krea2_override_accepts_single_module_native_lora(_raise_if_not 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. - 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 + 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") @@ -392,24 +395,94 @@ def test_automatic_probe_accepts_lokr_lora(_raise_if_not_file) -> None: 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_lora(_raise_if_not_file) -> None: +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( - _native_lokr_lora(), {**_REQUIRED_FIELDS, "base": BaseModelType.Krea2} + _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__` 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 = 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(), - } + 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}) 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 333acc61d92..b8ddbfdf7a9 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 @@ -11,10 +11,7 @@ KREA2_LORA_QWEN3VL_PREFIX, KREA2_LORA_TRANSFORMER_PREFIX, ) -from invokeai.backend.patches.lora_conversions.krea2_lora_conversion_utils import ( - is_state_dict_likely_krea2_lora, - lora_model_from_krea2_state_dict, -) +from invokeai.backend.patches.lora_conversions.krea2_lora_conversion_utils import 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, ) @@ -387,11 +384,13 @@ def test_unrecognized_kohya_flattened_keys_are_left_untouched(flat_module: str) # two groups and aborting the load on the orphaned `{'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. + # A weight-decomposed LoKr also saves a `dora_scale`. The converter drops it (there is no DoRA+LoKr + # layer), so the module still converts and applies as plain LoKr. + (("lokr_w1", "lokr_w2", "alpha", "dora_scale"), "transformer_blocks.6.attn.to_q"), + # `dora_scale` is deliberately not combined with LoHa or full here: it would orphan the same way, but + # even grouped correctly `any_lora_layer_from_state_dict` tests `dora_scale` before `hada_w1_a`, so + # the module 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( @@ -494,13 +493,15 @@ def test_native_krea2_top_level_linear_keys_are_remapped() -> None: assert expected_keys < set(model.layers) -def test_lokr_layer_produces_lokr_layer() -> None: +def test_lokr_layer_reconstructs_the_kronecker_product() -> 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. + # a lora_A/lora_B pair. They must survive _group_by_layer intact and land on the right LoKRLayer operand. + # The factors are random rather than all-ones so a swapped operand changes the result: + # kron(ones(2, 2), ones(3, 4)) == kron(ones(3, 4), ones(2, 2)). + w1, w2 = torch.randn(2, 2), torch.randn(3, 4) 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.lokr_w1": w1, + "transformer.text_fusion.0.attn.to_q.lokr_w2": w2, "transformer.text_fusion.0.attn.to_q.alpha": torch.tensor(1.0), } @@ -508,75 +509,80 @@ def test_lokr_layer_produces_lokr_layer() -> None: 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) + torch.testing.assert_close(layer.get_weight(torch.empty(6, 8)), torch.kron(w1, w2)) + # Unfactored LoKr has no rank, so alpha does not rescale it. + assert layer.scale() == 1.0 -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. +def test_factored_lokr_layer_reconstructs_the_kronecker_product() -> None: + # LoKr may factor either Kronecker operand further into an `_a`/`_b` pair. Each half has to reach its own + # operand, and alpha scales by the rank of the factorisation. + w1_a, w1_b = torch.randn(2, 1), torch.randn(1, 2) + w2_a, w2_b = torch.randn(3, 1), torch.randn(1, 4) 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), + "transformer.text_fusion.0.attn.to_q.lokr_w1_a": w1_a, + "transformer.text_fusion.0.attn.to_q.lokr_w1_b": w1_b, + "transformer.text_fusion.0.attn.to_q.lokr_w2_a": w2_a, + "transformer.text_fusion.0.attn.to_q.lokr_w2_b": w2_b, + "transformer.text_fusion.0.attn.to_q.alpha": torch.tensor(4.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.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 + torch.testing.assert_close(layer.get_weight(torch.empty(6, 8)), torch.kron(w1_a @ w1_b, w2_a @ w2_b)) + assert layer.scale() == 4.0 / w1_b.shape[0] 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. + w1, w2 = torch.randn(2, 2), torch.randn(3, 4) 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), + "diffusion_model.txtfusion.layerwise_blocks.0.attn.wq.lokr_w1": w1, + "diffusion_model.txtfusion.layerwise_blocks.0.attn.wq.lokr_w2": w2, + "diffusion_model.txtfusion.refiner_blocks.1.mlp.down.lokr_w1": w2, + "diffusion_model.txtfusion.refiner_blocks.1.mlp.down.lokr_w2": w1, } 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()) + to_q = model.layers[f"{KREA2_LORA_TRANSFORMER_PREFIX}text_fusion.layerwise_blocks.0.attn.to_q"] + down = model.layers[f"{KREA2_LORA_TRANSFORMER_PREFIX}text_fusion.refiner_blocks.1.ff.down"] + assert len(model.layers) == 2 + assert isinstance(to_q, LoKRLayer) and isinstance(down, LoKRLayer) + torch.testing.assert_close(to_q.get_weight(torch.empty(6, 8)), torch.kron(w1, w2)) + torch.testing.assert_close(down.get_weight(torch.empty(6, 8)), torch.kron(w2, w1)) -def test_is_state_dict_likely_krea2_lora_accepts_lokr() -> None: +@pytest.mark.parametrize("magnitude_suffix", ["dora_scale", "magnitude"]) +def test_lokr_layer_drops_a_dora_magnitude(magnitude_suffix: str) -> None: + # A weight-decomposed LoKr also saves a DoRA magnitude. Kept, it would dispatch the layer to DoRALayer, + # which fails on the missing lora_up/lora_down. There is no DoRA+LoKr layer, so it applies as plain LoKr. + w1, w2 = torch.randn(2, 2), torch.randn(3, 4) 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), + "transformer.text_fusion.0.attn.to_q.lokr_w1": w1, + "transformer.text_fusion.0.attn.to_q.lokr_w2": w2, + f"transformer.text_fusion.0.attn.to_q.{magnitude_suffix}": torch.ones(8), } - 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), - } + model = lora_model_from_krea2_state_dict(state_dict) - assert not is_state_dict_likely_krea2_lora(state_dict) + layer = model.layers[f"{KREA2_LORA_TRANSFORMER_PREFIX}text_fusion.0.attn.to_q"] + assert isinstance(layer, LoKRLayer) + torch.testing.assert_close(layer.get_weight(torch.empty(6, 8)), torch.kron(w1, w2)) 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. + w1, w2_a, w2_b = torch.randn(2, 2), torch.randn(3, 1), torch.randn(1, 4) 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.lokr_w1": w1, + "lora_unet_txtfusion_layerwise_blocks_0_attn_wq.lokr_w2_a": w2_a, + "lora_unet_txtfusion_layerwise_blocks_0_attn_wq.lokr_w2_b": w2_b, "lora_unet_txtfusion_layerwise_blocks_0_attn_wq.alpha": torch.tensor(4.0), } @@ -584,4 +590,5 @@ def test_kohya_flattened_lokr_converts_onto_its_real_module() -> None: 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) + torch.testing.assert_close(layer.get_weight(torch.empty(6, 8)), torch.kron(w1, w2_a @ w2_b)) + assert layer.scale() == 4.0 / w2_b.shape[0]