diff --git a/invokeai/backend/model_manager/configs/lora.py b/invokeai/backend/model_manager/configs/lora.py index bbdd0f46561..7fccf20184a 100644 --- a/invokeai/backend/model_manager/configs/lora.py +++ b/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. 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 029ab6780a6..d5d3e678ea3 100644 --- a/invokeai/backend/patches/lora_conversions/krea2_lora_conversion_utils.py +++ b/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 ``.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", } 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 28499f2936d..25757344aaf 100644 --- a/tests/backend/model_manager/configs/test_krea2_lora_config.py +++ b/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}) 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 5fc9b4d2d81..7e7fd48b352 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 @@ -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,17 +338,19 @@ 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 @@ -352,7 +358,9 @@ def test_unrecognized_kohya_flattened_keys_are_left_untouched(flat_module: str) # 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)