diff --git a/invokeai/backend/model_manager/configs/lora.py b/invokeai/backend/model_manager/configs/lora.py index bbdd0f46561..28c3ca01e8f 100644 --- a/invokeai/backend/model_manager/configs/lora.py +++ b/invokeai/backend/model_manager/configs/lora.py @@ -1,5 +1,6 @@ import re from abc import ABC +from collections.abc import Callable from pathlib import Path from typing import ( Any, @@ -47,6 +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.wan_lora_constants import ( detect_wan_lora_variant, has_non_wan_architecture_keys, @@ -908,12 +910,15 @@ def _has_krea2_lora_keys(state_dict: dict[str | int, Any]) -> bool: } -def _lora_weight_keys_are_all_paired(state_dict: dict[str | int, Any], prefixes: tuple[str, ...] | None = None) -> bool: - """True if *every* lora_A/lora_B/lora_down/lora_up weight (optionally restricted to `prefixes`) has its - partner half present. Returns True when there are no such weights at all (nothing to invalidate).""" +def _lora_weight_keys_are_all_paired( + state_dict: dict[str | int, Any], key_filter: Callable[[str], bool] | None = None +) -> bool: + """True if *every* lora_A/lora_B/lora_down/lora_up weight (optionally restricted to the keys `key_filter` + accepts) has its partner half present. Returns True when there are no such weights at all (nothing to + invalidate).""" string_keys = {key for key in state_dict if isinstance(key, str)} for key in string_keys: - if prefixes is not None and not key.startswith(prefixes): + if key_filter is not None and not key_filter(key): continue for suffix, partner_suffix in _LORA_PAIR_PARTNERS.items(): if key.endswith(suffix): @@ -923,8 +928,9 @@ def _lora_weight_keys_are_all_paired(state_dict: dict[str | int, Any], prefixes: return True -def _has_complete_lora_pair(state_dict: dict[str | int, Any], prefixes: tuple[str, ...] | None = None) -> bool: - """True if at least one complete lora_A/B (or lora_down/up) pair exists, optionally under `prefixes`. +def _has_complete_lora_pair(state_dict: dict[str | int, Any], key_filter: Callable[[str], bool] | None = None) -> bool: + """True if at least one complete lora_A/B (or lora_down/up) pair exists, optionally restricted to the + keys `key_filter` accepts. Note this only requires a *complete* pair to exist; it does not by itself reject dangling halves elsewhere — callers pair it with :func:`_lora_weight_keys_are_all_paired` (over the whole state dict) @@ -932,7 +938,7 @@ def _has_complete_lora_pair(state_dict: dict[str | int, Any], prefixes: tuple[st """ string_keys = {key for key in state_dict if isinstance(key, str)} for key in string_keys: - if prefixes is not None and not key.startswith(prefixes): + if key_filter is not None and not key_filter(key): continue for suffix, partner_suffix in _LORA_PAIR_PARTNERS.items(): if key.endswith(suffix) and f"{key[: -len(suffix)]}{partner_suffix}" in string_keys: @@ -940,8 +946,9 @@ def _has_complete_lora_pair(state_dict: dict[str | int, Any], prefixes: tuple[st 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). +# 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`. _KREA2_SUPPORTED_LORA_PREFIXES = ( "transformer.transformer_blocks.", "transformer_blocks.", @@ -955,22 +962,26 @@ def _has_complete_lora_pair(state_dict: dict[str | int, Any], prefixes: tuple[st "diffusion_model.tproj.", "diffusion_model.txtmlp.", "diffusion_model.last.linear.", - # kohya/LyCORIS flattens that same native layout into `lora_unet_` (see - # krea2_lora_conversion_utils._maybe_convert_kohya_krea2_state_dict). Spelled out per top-level module - # rather than as a bare `lora_unet_`, which would match every other architecture's kohya LoRA too. - "lora_unet_blocks_", - "lora_unet_txtfusion_", - "lora_unet_first.", - "lora_unet_tmlp_", - "lora_unet_tproj_", - "lora_unet_txtmlp_", - "lora_unet_last_linear.", "base_model.model.transformer.transformer_blocks.", "text_encoder.", "base_model.model.text_encoder.", ) +def _key_is_supported_krea2_layout(key: str) -> bool: + """True if `key` names a module in a layout the Krea-2 converter can map. + + The dotted layouts are matched by prefix. The kohya/LyCORIS layout is not, because its `lora_unet_` + spelling is shared: Wan writes `lora_unet_blocks__...` and Anima `lora_unet_[llm_adapter_]blocks_ + _...`, so per-module prefixes such as `lora_unet_blocks_` still sweep those files into the explicit + Krea-2 override, where they install and then silently no-op at generation time. Asking the converter's + own un-flattener whether the flattened path reconstructs to a Krea-2 module leaf rejects them, and as a + bonus accepts the doubled-separator spelling (`lora_unet__blocks_...`) that the converter tolerates but + no prefix spelled out (review 4888833569, notes 2 and 3). + """ + return key.startswith(_KREA2_SUPPORTED_LORA_PREFIXES) or is_kohya_krea2_lora_key(key) + + class LoRA_LyCORIS_Krea2_Config(LoRA_LyCORIS_Config_Base, Config_Base): """Model config for Krea-2 LoRA models in LyCORIS (single-file diffusers PEFT) format.""" @@ -983,7 +994,7 @@ 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, _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): diff --git a/invokeai/backend/patches/lora_conversions/krea2_lora_constants.py b/invokeai/backend/patches/lora_conversions/krea2_lora_constants.py index 07dfa12e36a..433a4d1242d 100644 --- a/invokeai/backend/patches/lora_conversions/krea2_lora_constants.py +++ b/invokeai/backend/patches/lora_conversions/krea2_lora_constants.py @@ -1,8 +1,125 @@ -# Krea-2 LoRA prefix constants. -# These prefixes namespace LoRA patch keys when applying them to Krea-2 models. +"""Krea-2 LoRA prefix constants and kohya key-reconstruction helpers. + +The prefixes namespace LoRA patch keys when applying them to Krea-2 models. + +The kohya/LyCORIS reconstruction helpers live here rather than in ``krea2_lora_conversion_utils`` so that +``model_manager.configs.lora`` can identify a flattened Krea-2 LoRA without importing the converter: the +converter pulls in the patch layers, which import ``model_manager.load``, closing an import cycle back into +``model_manager.configs``. Same reason ``anima_lora_constants`` exists. +""" + +from invokeai.backend.patches.lora_conversions.kohya_key_utils import ( + INDEX_PLACEHOLDER, + ParsingTree, + insert_periods_into_kohya_key, +) # Prefix for Krea-2 transformer (Krea2Transformer2DModel) LoRA layers. KREA2_LORA_TRANSFORMER_PREFIX = "lora_transformer-" # Prefix for Krea-2 Qwen3-VL text encoder LoRA layers. KREA2_LORA_QWEN3VL_PREFIX = "lora_qwen3vl-" + + +# --- Kohya / LyCORIS (flattened) -> native key mapping --------------------------------------------------------- +# sd-scripts and LyCORIS flatten the module path (``path.replace(".", "_")``) and prefix it with +# ``lora_unet_``, e.g. ``lora_unet_blocks_6_attn_wv.lora_down.weight``. Flattening is lossy — nothing in the key +# records where a '_' used to be a '.' — so we reconstruct the dotted path against the native module vocabulary +# below and accept it only if it lands on a leaf. A key we cannot reconstruct with certainty is left untouched +# rather than rewritten into a plausible-looking key that matches no module. +KREA2_KOHYA_PREFIX = "lora_unet_" + +# Native Krea-2 transformer/text-fusion block leaves. Only the Linears are listed: the non-Linear natives +# (``mod.lin``, ``prenorm``/``postnorm``, ``attn.qknorm.*``, ``last.norm``/``last.modulation``) have no Linear +# counterpart in the diffusers layout — ``mod.lin`` for instance is folded into the ``scale_shift_table`` +# parameter — so an adapter targeting them cannot be applied, and renaming it anyway would turn "unsupported" +# into a silent no-op. +_NATIVE_KREA2_BLOCK_SUBTREE: ParsingTree = { + "attn": {"wq": {}, "wk": {}, "wv": {}, "wo": {}, "gate": {}}, + "mlp": {"gate": {}, "up": {}, "down": {}}, +} + +# Parsing tree for the native (ComfyUI) Krea-2 module layout, i.e. the keys the renames in +# ``krea2_lora_conversion_utils`` understand. Walking it resolves the flattened form's only real ambiguity — ``layerwise_blocks`` / ``refiner_blocks`` are +# the native components that themselves contain an underscore. +_KREA2_NATIVE_KOHYA_PARSING_TREE: ParsingTree = { + "blocks": {INDEX_PLACEHOLDER: _NATIVE_KREA2_BLOCK_SUBTREE}, + "txtfusion": { + "layerwise_blocks": {INDEX_PLACEHOLDER: _NATIVE_KREA2_BLOCK_SUBTREE}, + "refiner_blocks": {INDEX_PLACEHOLDER: _NATIVE_KREA2_BLOCK_SUBTREE}, + "projector": {}, + }, + "first": {}, + # Literal indices rather than INDEX_PLACEHOLDER: these are ``nn.Sequential`` stages, and only the + # positions listed in the converter's ``_NATIVE_KREA2_TOP_LEVEL_RENAMES`` hold a Linear — the rest are + # activations with no weights. Accepting any index would rewrite e.g. ``lora_unet_tmlp_1`` into ``tmlp.1.*``, which the + # native pass then does not recognize, leaving a half-converted key instead of the untouched original. + "tmlp": {"0": {}, "2": {}}, + "tproj": {"1": {}}, + "txtmlp": {"1": {}, "3": {}}, + "last": {"linear": {}}, +} + + +def _kohya_module_path_is_leaf(module_path: str, parsing_tree: ParsingTree) -> bool: + """True if a dotted module path walks the tree all the way to a leaf. + + ``insert_periods_into_kohya_key`` only rejects *leftover* tokens, so a prefix of a real path (e.g. + ``blocks.0.attn``) parses cleanly without naming a module. Requiring a leaf rejects those. + """ + subtree = parsing_tree + for component in module_path.split("."): + # Mirror ``insert_periods_into_kohya_key``'s precedence: an exact match wins over the index + # placeholder. Without that, a numeric component would always be looked up as INDEX_PLACEHOLDER and + # a tree enumerating the specific indices it accepts (``tmlp`` below) could never reach its leaves. + if component in subtree: + subtree = subtree[component] + elif component.isnumeric() and INDEX_PLACEHOLDER in subtree: + subtree = subtree[INDEX_PLACEHOLDER] + else: + return False + return not subtree + + +def unflatten_kohya_krea2_module_path(flat_path: str) -> str | None: + """Reconstruct a dotted native Krea-2 module path from its kohya-flattened form. + + Returns ``None`` when the reconstruction is not a native module path the converter can map, in which case + the caller must leave the key alone. + """ + try: + module_path = insert_periods_into_kohya_key(flat_path, _KREA2_NATIVE_KOHYA_PARSING_TREE) + except ValueError: + # Tokens left over: not a native Krea-2 module path. + return None + return module_path if _kohya_module_path_is_leaf(module_path, _KREA2_NATIVE_KOHYA_PARSING_TREE) else None + + +def split_kohya_krea2_key(key: str | int) -> tuple[str, str, str] | None: + """Split a kohya/LyCORIS key into (flattened module path, separator, weight suffix). + + Returns ``None`` for anything not in the kohya layout, including the non-string keys that ``.pt`` / + ``.ckpt`` sources can carry. The flattened module path runs up to the first '.'; the weight suffix + (``lora_down.weight``, ``alpha``, ...) follows it. Some writers emit a doubled separator after the + prefix, hence the ``lstrip``. Every caller splits through here so the converter's per-module gate, the + rewrite it guards, and identification can never disagree about which module a key belongs to. + """ + if not isinstance(key, str) or not key.startswith(KREA2_KOHYA_PREFIX): + return None + flat_path, dot, weight_suffix = key[len(KREA2_KOHYA_PREFIX) :].lstrip("_").partition(".") + return flat_path, dot, weight_suffix + + +def is_kohya_krea2_lora_key(key: str | int) -> bool: + """True if ``key`` is a kohya-flattened key naming a Krea-2 module the converter can reconstruct. + + Identification (``model_manager.configs.lora``) uses this instead of matching ``lora_unet_`` + prefixes. That spelling is not Krea-2's alone: Wan writes ``lora_unet_blocks__...`` and Anima + ``lora_unet_[llm_adapter_]blocks__...``, so a prefix match sweeps their kohya LoRAs into the + explicit Krea-2 override, where they install and then silently no-op at generation time. Reconstructing + the path against the native Krea-2 module vocabulary rejects them - ``self_attn``, ``cross_attn`` and + ``mlp_layer0`` are not leaves in it - while still accepting the doubled-separator spelling that the + converter tolerates but no ``lora_unet_`` prefix spells out. + """ + split = split_kohya_krea2_key(key) + return split is not None and unflatten_kohya_krea2_module_path(split[0]) is not None 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..f89154b3bb9 100644 --- a/invokeai/backend/patches/lora_conversions/krea2_lora_conversion_utils.py +++ b/invokeai/backend/patches/lora_conversions/krea2_lora_conversion_utils.py @@ -17,14 +17,11 @@ from invokeai.backend.patches.layers.base_layer_patch import BaseLayerPatch from invokeai.backend.patches.layers.utils import any_lora_layer_from_state_dict -from invokeai.backend.patches.lora_conversions.kohya_key_utils import ( - INDEX_PLACEHOLDER, - ParsingTree, - insert_periods_into_kohya_key, -) from invokeai.backend.patches.lora_conversions.krea2_lora_constants import ( KREA2_LORA_QWEN3VL_PREFIX, KREA2_LORA_TRANSFORMER_PREFIX, + split_kohya_krea2_key, + unflatten_kohya_krea2_module_path, ) from invokeai.backend.patches.model_patch_raw import ModelPatchRaw @@ -115,80 +112,6 @@ def _maybe_convert_native_krea2_state_dict( return converted_state_dict -# --- Kohya / LyCORIS (flattened) -> native key mapping --------------------------------------------------------- -# sd-scripts and LyCORIS flatten the module path (``path.replace(".", "_")``) and prefix it with -# ``lora_unet_``, e.g. ``lora_unet_blocks_6_attn_wv.lora_down.weight``. Flattening is lossy — nothing in the key -# records where a '_' used to be a '.' — so we reconstruct the dotted path against the native module vocabulary -# below and accept it only if it lands on a leaf. A key we cannot reconstruct with certainty is left untouched -# rather than rewritten into a plausible-looking key that matches no module. -_KREA2_KOHYA_PREFIX = "lora_unet_" - -# Native Krea-2 transformer/text-fusion block leaves. Only the Linears are listed: the non-Linear natives -# (``mod.lin``, ``prenorm``/``postnorm``, ``attn.qknorm.*``, ``last.norm``/``last.modulation``) have no Linear -# counterpart in the diffusers layout — ``mod.lin`` for instance is folded into the ``scale_shift_table`` -# parameter — so an adapter targeting them cannot be applied, and renaming it anyway would turn "unsupported" -# into a silent no-op. -_NATIVE_KREA2_BLOCK_SUBTREE: ParsingTree = { - "attn": {"wq": {}, "wk": {}, "wv": {}, "wo": {}, "gate": {}}, - "mlp": {"gate": {}, "up": {}, "down": {}}, -} - -# Parsing tree for the native (ComfyUI) Krea-2 module layout, i.e. the keys the renames above understand. -# Walking it resolves the flattened form's only real ambiguity — ``layerwise_blocks`` / ``refiner_blocks`` are -# the native components that themselves contain an underscore. -_KREA2_NATIVE_KOHYA_PARSING_TREE: ParsingTree = { - "blocks": {INDEX_PLACEHOLDER: _NATIVE_KREA2_BLOCK_SUBTREE}, - "txtfusion": { - "layerwise_blocks": {INDEX_PLACEHOLDER: _NATIVE_KREA2_BLOCK_SUBTREE}, - "refiner_blocks": {INDEX_PLACEHOLDER: _NATIVE_KREA2_BLOCK_SUBTREE}, - "projector": {}, - }, - "first": {}, - # Literal indices rather than INDEX_PLACEHOLDER: these are ``nn.Sequential`` stages, and only the - # positions listed in ``_NATIVE_KREA2_TOP_LEVEL_RENAMES`` hold a Linear — the rest are activations with - # no weights. Accepting any index would rewrite e.g. ``lora_unet_tmlp_1`` into ``tmlp.1.*``, which the - # native pass then does not recognize, leaving a half-converted key instead of the untouched original. - "tmlp": {"0": {}, "2": {}}, - "tproj": {"1": {}}, - "txtmlp": {"1": {}, "3": {}}, - "last": {"linear": {}}, -} - - -def _kohya_module_path_is_leaf(module_path: str, parsing_tree: ParsingTree) -> bool: - """True if a dotted module path walks the tree all the way to a leaf. - - ``insert_periods_into_kohya_key`` only rejects *leftover* tokens, so a prefix of a real path (e.g. - ``blocks.0.attn``) parses cleanly without naming a module. Requiring a leaf rejects those. - """ - subtree = parsing_tree - for component in module_path.split("."): - # Mirror ``insert_periods_into_kohya_key``'s precedence: an exact match wins over the index - # placeholder. Without that, a numeric component would always be looked up as INDEX_PLACEHOLDER and - # a tree enumerating the specific indices it accepts (``tmlp`` below) could never reach its leaves. - if component in subtree: - subtree = subtree[component] - elif component.isnumeric() and INDEX_PLACEHOLDER in subtree: - subtree = subtree[INDEX_PLACEHOLDER] - else: - return False - return not subtree - - -def _unflatten_kohya_krea2_module_path(flat_path: str) -> str | None: - """Reconstruct a dotted native Krea-2 module path from its kohya-flattened form. - - Returns ``None`` when the reconstruction is not a native module path this converter can map, in which case - the caller must leave the key alone. - """ - try: - module_path = insert_periods_into_kohya_key(flat_path, _KREA2_NATIVE_KOHYA_PARSING_TREE) - except ValueError: - # Tokens left over: not a native Krea-2 module path. - return None - return module_path if _kohya_module_path_is_leaf(module_path, _KREA2_NATIVE_KOHYA_PARSING_TREE) else None - - def _maybe_convert_kohya_krea2_state_dict( state_dict: Dict[str, torch.Tensor], ) -> Dict[str, torch.Tensor]: @@ -202,8 +125,9 @@ def _maybe_convert_kohya_krea2_state_dict( # So collect each flattened module's suffixes first and convert only modules where *all* of them convert. suffixes_by_flat_path: dict[str, set[str]] = {} for key in state_dict: - if isinstance(key, str) and key.startswith(_KREA2_KOHYA_PREFIX): - flat_path, _, weight_suffix = key[len(_KREA2_KOHYA_PREFIX) :].lstrip("_").partition(".") + split = split_kohya_krea2_key(key) + if split is not None: + flat_path, _, weight_suffix = split suffixes_by_flat_path.setdefault(flat_path, set()).add(f".{weight_suffix}") fully_convertible_flat_paths = { flat_path @@ -215,11 +139,10 @@ def _maybe_convert_kohya_krea2_state_dict( source_keys: dict[str, str] = {} for key, value in state_dict.items(): converted_key = key - if isinstance(key, str) and key.startswith(_KREA2_KOHYA_PREFIX): - # The flattened module path runs up to the first '.'; the weight suffix (``lora_down.weight``, - # ``alpha``, ...) follows it. Some writers emit a doubled separator after the prefix. - flat_path, dot, weight_suffix = key[len(_KREA2_KOHYA_PREFIX) :].lstrip("_").partition(".") - module_path = _unflatten_kohya_krea2_module_path(flat_path) + split = split_kohya_krea2_key(key) + if split is not None: + flat_path, dot, weight_suffix = split + module_path = unflatten_kohya_krea2_module_path(flat_path) # 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 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..b10064d887e 100644 --- a/tests/backend/model_manager/configs/test_krea2_lora_config.py +++ b/tests/backend/model_manager/configs/test_krea2_lora_config.py @@ -70,7 +70,7 @@ def test_explicit_krea2_override_accepts_kohya_transformer_only_lora(_raise_if_n @patch("invokeai.backend.model_manager.configs.lora.raise_if_not_file") def test_explicit_krea2_override_still_rejects_a_foreign_kohya_lora(_raise_if_not_file) -> None: - """The new `lora_unet_` entries are per-module, so another architecture's kohya LoRA is not swept in.""" + """A kohya key whose flattened path does not reconstruct to a Krea-2 module is not swept in.""" mod = MagicMock() mod.load_state_dict.return_value = { "lora_unet_double_blocks_0_img_attn_proj.lora_down.weight": object(), @@ -81,6 +81,104 @@ def test_explicit_krea2_override_still_rejects_a_foreign_kohya_lora(_raise_if_no LoRA_LyCORIS_Krea2_Config.from_model_on_disk(mod, {**_REQUIRED_FIELDS, "base": BaseModelType.Krea2}) +# Wan and Anima flatten their kohya module paths under the very same `lora_unet_blocks__` spelling as +# Krea-2 (wan_lora_conversion_utils._KOHYA_KEY_REGEX, anima_lora_constants._KOHYA_ANIMA_RE). Matching that as +# a prefix accepted a mislabeled Wan/Anima file under the explicit Krea-2 override, where it installed and +# then silently no-op'd at generation time - every layer warn-skipped, because the un-flattener rejects +# `self_attn`/`cross_attn`/`mlp_layer0` (review 4888833569, note 2). +_FOREIGN_KOHYA_MODULES = [ + "blocks_0_self_attn_q", # Wan + "blocks_0_cross_attn_k", # Wan + "blocks_0_ffn_0", # Wan + "blocks_0_mlp_layer0", # Anima + "blocks_0_adaln_modulation_1", # Anima + "llm_adapter_blocks_0_self_attn_q_proj", # Anima +] + + +@pytest.mark.parametrize("kohya_module", _FOREIGN_KOHYA_MODULES) +@patch("invokeai.backend.model_manager.configs.lora.raise_if_not_file") +def test_explicit_krea2_override_rejects_foreign_lora_unet_blocks_lora(_raise_if_not_file, kohya_module: str) -> None: + mod = MagicMock() + mod.load_state_dict.return_value = { + f"lora_unet_{kohya_module}.lora_down.weight": object(), + f"lora_unet_{kohya_module}.lora_up.weight": object(), + } + + with pytest.raises(NotAMatchError): + LoRA_LyCORIS_Krea2_Config.from_model_on_disk(mod, {**_REQUIRED_FIELDS, "base": BaseModelType.Krea2}) + + +# The converter deliberately tolerates the doubled separator some writers emit +# (test_kohya_flattened_krea2_keys_tolerate_doubled_separator), but no `lora_unet_` prefix spelled it out, so +# this transformer-only adapter could not be installed at all (review 4888833569, note 3). +@pytest.mark.parametrize( + "kohya_module", + [ + "blocks_0_attn_wq", + "txtfusion_refiner_blocks_1_mlp_up", + "tproj_1", + "last_linear", + ], +) +@patch("invokeai.backend.model_manager.configs.lora.raise_if_not_file") +def test_explicit_krea2_override_accepts_doubled_separator_kohya_lora(_raise_if_not_file, kohya_module: str) -> None: + mod = MagicMock() + mod.load_state_dict.return_value = { + f"lora_unet__{kohya_module}.lora_down.weight": object(), + f"lora_unet__{kohya_module}.lora_up.weight": object(), + } + + config = LoRA_LyCORIS_Krea2_Config.from_model_on_disk(mod, {**_REQUIRED_FIELDS, "base": BaseModelType.Krea2}) + + assert config.base is BaseModelType.Krea2 + + +# The same module vocabulary in the single-separator spelling, to pin that the un-flattener - not a prefix +# list - is what admits a kohya adapter now. +@pytest.mark.parametrize( + "kohya_module", + [ + "blocks_0_attn_wq", + "blocks_0_mlp_down", + "txtfusion_layerwise_blocks_2_attn_gate", + "txtfusion_projector", + "first", + "tmlp_0", + "tproj_1", + "txtmlp_3", + "last_linear", + ], +) +@patch("invokeai.backend.model_manager.configs.lora.raise_if_not_file") +def test_explicit_krea2_override_accepts_every_kohya_module(_raise_if_not_file, kohya_module: str) -> None: + mod = MagicMock() + mod.load_state_dict.return_value = { + f"lora_unet_{kohya_module}.lora_down.weight": object(), + f"lora_unet_{kohya_module}.lora_up.weight": object(), + } + + config = LoRA_LyCORIS_Krea2_Config.from_model_on_disk(mod, {**_REQUIRED_FIELDS, "base": BaseModelType.Krea2}) + + assert config.base is BaseModelType.Krea2 + + +# `tmlp`/`tproj`/`txtmlp` are nn.Sequential stages where only some positions hold a Linear. The parsing tree +# lists those positions literally, so an adapter on an activation index is not a Krea-2 module and must not +# install - the converter would leave its keys untouched and the layer would warn-skip. +@pytest.mark.parametrize("kohya_module", ["tmlp_1", "tproj_0", "txtmlp_2"]) +@patch("invokeai.backend.model_manager.configs.lora.raise_if_not_file") +def test_explicit_krea2_override_rejects_non_linear_sequential_index(_raise_if_not_file, kohya_module: str) -> None: + mod = MagicMock() + mod.load_state_dict.return_value = { + f"lora_unet_{kohya_module}.lora_down.weight": object(), + f"lora_unet_{kohya_module}.lora_up.weight": object(), + } + + with pytest.raises(NotAMatchError): + LoRA_LyCORIS_Krea2_Config.from_model_on_disk(mod, {**_REQUIRED_FIELDS, "base": BaseModelType.Krea2}) + + @patch("invokeai.backend.model_manager.configs.lora.raise_if_not_file") def test_explicit_krea2_override_accepts_ambiguous_transformer_only_lora(_raise_if_not_file) -> None: config = LoRA_LyCORIS_Krea2_Config.from_model_on_disk(