diff --git a/invokeai/backend/model_manager/configs/lora.py b/invokeai/backend/model_manager/configs/lora.py index fbf7cfa8b6c..f6f97eba87a 100644 --- a/invokeai/backend/model_manager/configs/lora.py +++ b/invokeai/backend/model_manager/configs/lora.py @@ -954,6 +954,16 @@ 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.", 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..ce364fee229 100644 --- a/invokeai/backend/patches/lora_conversions/krea2_lora_conversion_utils.py +++ b/invokeai/backend/patches/lora_conversions/krea2_lora_conversion_utils.py @@ -5,6 +5,9 @@ ``transformer..lora_A.weight`` / ``lora_B.weight``. The distinctive Krea-2 module is the ``text_fusion`` stage, which we use to disambiguate from Qwen-Image / Z-Image LoRAs (which otherwise share the ``transformer.transformer_blocks.`` prefix). + +Two other key layouts are normalized onto that one before conversion: the native (ComfyUI) module naming, +and the kohya / LyCORIS layout that additionally flattens the module path into ``lora_unet_``. """ import re @@ -14,6 +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, @@ -94,7 +102,10 @@ def _maybe_convert_native_krea2_state_dict( return state_dict converted_state_dict: Dict[str, torch.Tensor] = {} for key, value in state_dict.items(): - converted_key = _native_krea2_key_to_diffusers(key) if _looks_like_native_krea2_key(key) else key + # `.pt`/`.ckpt` sources can carry non-string keys. They are never native Krea-2 keys, but the + # substring tests in `_looks_like_native_krea2_key` raise TypeError rather than returning False. + is_native = isinstance(key, str) and _looks_like_native_krea2_key(key) + converted_key = _native_krea2_key_to_diffusers(key) if is_native else key if converted_key in converted_state_dict: raise ValueError( f"Krea-2 LoRA has conflicting layers that normalize to the same target '{converted_key}'. " @@ -104,6 +115,112 @@ 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]: + """Rewrite kohya/LyCORIS flattened Krea-2 keys to the dotted native layout, leaving all others untouched.""" + converted_state_dict: Dict[str, torch.Tensor] = {} + 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) + # 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. + if module_path is not None and f".{weight_suffix}" in _SUFFIX_TO_VALUE_KEY: + converted_key = f"{module_path}{dot}{weight_suffix}" + if converted_key in converted_state_dict: + raise ValueError( + f"Krea-2 LoRA has conflicting layers that normalize to the same target '{converted_key}' " + f"(from '{source_keys[converted_key]}' and '{key}'). This mixed layout is unsupported - " + "refusing to silently drop one of the layers." + ) + converted_state_dict[converted_key] = value + source_keys[converted_key] = str(key) + return converted_state_dict + + def is_state_dict_likely_krea2_lora(state_dict: dict[str | int, torch.Tensor]) -> bool: """Checks if the provided state dict is likely a Krea-2 LoRA. @@ -125,7 +242,9 @@ def lora_model_from_krea2_state_dict(state_dict: Dict[str, torch.Tensor], alpha: as ``alpha=rank`` internally (the common diffusers default). """ layers: dict[str, BaseLayerPatch] = {} - # Normalize native (ComfyUI) naming to the diffusers layout so the rest of the converter is layout-agnostic. + # Normalize the kohya/LyCORIS flattened naming (``lora_unet_blocks_6_attn_wv``) to the dotted native layout, + # then the native (ComfyUI) naming to the diffusers layout, so the rest of the converter is layout-agnostic. + state_dict = _maybe_convert_kohya_krea2_state_dict(state_dict) state_dict = _maybe_convert_native_krea2_state_dict(state_dict) grouped_state_dict = _group_by_layer(state_dict) 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..28499f2936d 100644 --- a/tests/backend/model_manager/configs/test_krea2_lora_config.py +++ b/tests/backend/model_manager/configs/test_krea2_lora_config.py @@ -45,6 +45,42 @@ def _diffusion_model_transformer_only_lora() -> MagicMock: return mod +def _kohya_transformer_only_lora() -> MagicMock: + mod = MagicMock() + mod.load_state_dict.return_value = { + "lora_unet_blocks_0_attn_wq.lora_down.weight": object(), + "lora_unet_blocks_0_attn_wq.lora_up.weight": object(), + } + return mod + + +@patch("invokeai.backend.model_manager.configs.lora.raise_if_not_file") +def test_explicit_krea2_override_accepts_kohya_transformer_only_lora(_raise_if_not_file) -> None: + """The converter understands the flattened kohya layout, so the override escape hatch must too. + + Auto-detection can't reach a transformer-only kohya adapter: it has no `txtfusion` key, and the + gated-attention fallback looks for the dotted `.attn.wq.` spelling that flattening destroys. + """ + config = LoRA_LyCORIS_Krea2_Config.from_model_on_disk( + _kohya_transformer_only_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_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.""" + mod = MagicMock() + mod.load_state_dict.return_value = { + "lora_unet_double_blocks_0_img_attn_proj.lora_down.weight": object(), + "lora_unet_double_blocks_0_img_attn_proj.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( diff --git a/tests/backend/patches/lora_conversions/lora_state_dicts/krea2_lora_kohya_format.py b/tests/backend/patches/lora_conversions/lora_state_dicts/krea2_lora_kohya_format.py new file mode 100644 index 00000000000..837e5ba071a --- /dev/null +++ b/tests/backend/patches/lora_conversions/lora_state_dicts/krea2_lora_kohya_format.py @@ -0,0 +1,133 @@ +"""Representative key layout of a kohya (sd-scripts) Krea-2 LoRA. + +Captured from `Krea2_Don-Martin_LoRA-step00001200.safetensors` (metadata: `ss_base_model_version=krea2`, +`ss_network_dim=32`, `ss_network_alpha=32`). The full adapter has 264 modules / 792 tensors covering all 28 +transformer blocks and both text-fusion stages; this fixture keeps the first and last transformer block, one +layerwise and one refiner text-fusion block, and every top-level module. + +kohya flattens the module path (`path.replace(".", "_")`) and prefixes it with `lora_unet_`, over the native +(ComfyUI) module names - so `lora_unet_blocks_6_attn_wv` denotes `transformer_blocks.6.attn.to_v`. +""" + +state_dict_keys: dict[str, list[int]] = { + "lora_unet_blocks_0_attn_gate.alpha": [], + "lora_unet_blocks_0_attn_gate.lora_down.weight": [32, 6144], + "lora_unet_blocks_0_attn_gate.lora_up.weight": [6144, 32], + "lora_unet_blocks_0_attn_wk.alpha": [], + "lora_unet_blocks_0_attn_wk.lora_down.weight": [32, 6144], + "lora_unet_blocks_0_attn_wk.lora_up.weight": [1536, 32], + "lora_unet_blocks_0_attn_wo.alpha": [], + "lora_unet_blocks_0_attn_wo.lora_down.weight": [32, 6144], + "lora_unet_blocks_0_attn_wo.lora_up.weight": [6144, 32], + "lora_unet_blocks_0_attn_wq.alpha": [], + "lora_unet_blocks_0_attn_wq.lora_down.weight": [32, 6144], + "lora_unet_blocks_0_attn_wq.lora_up.weight": [6144, 32], + "lora_unet_blocks_0_attn_wv.alpha": [], + "lora_unet_blocks_0_attn_wv.lora_down.weight": [32, 6144], + "lora_unet_blocks_0_attn_wv.lora_up.weight": [1536, 32], + "lora_unet_blocks_0_mlp_down.alpha": [], + "lora_unet_blocks_0_mlp_down.lora_down.weight": [32, 16384], + "lora_unet_blocks_0_mlp_down.lora_up.weight": [6144, 32], + "lora_unet_blocks_0_mlp_gate.alpha": [], + "lora_unet_blocks_0_mlp_gate.lora_down.weight": [32, 6144], + "lora_unet_blocks_0_mlp_gate.lora_up.weight": [16384, 32], + "lora_unet_blocks_0_mlp_up.alpha": [], + "lora_unet_blocks_0_mlp_up.lora_down.weight": [32, 6144], + "lora_unet_blocks_0_mlp_up.lora_up.weight": [16384, 32], + "lora_unet_blocks_27_attn_gate.alpha": [], + "lora_unet_blocks_27_attn_gate.lora_down.weight": [32, 6144], + "lora_unet_blocks_27_attn_gate.lora_up.weight": [6144, 32], + "lora_unet_blocks_27_attn_wk.alpha": [], + "lora_unet_blocks_27_attn_wk.lora_down.weight": [32, 6144], + "lora_unet_blocks_27_attn_wk.lora_up.weight": [1536, 32], + "lora_unet_blocks_27_attn_wo.alpha": [], + "lora_unet_blocks_27_attn_wo.lora_down.weight": [32, 6144], + "lora_unet_blocks_27_attn_wo.lora_up.weight": [6144, 32], + "lora_unet_blocks_27_attn_wq.alpha": [], + "lora_unet_blocks_27_attn_wq.lora_down.weight": [32, 6144], + "lora_unet_blocks_27_attn_wq.lora_up.weight": [6144, 32], + "lora_unet_blocks_27_attn_wv.alpha": [], + "lora_unet_blocks_27_attn_wv.lora_down.weight": [32, 6144], + "lora_unet_blocks_27_attn_wv.lora_up.weight": [1536, 32], + "lora_unet_blocks_27_mlp_down.alpha": [], + "lora_unet_blocks_27_mlp_down.lora_down.weight": [32, 16384], + "lora_unet_blocks_27_mlp_down.lora_up.weight": [6144, 32], + "lora_unet_blocks_27_mlp_gate.alpha": [], + "lora_unet_blocks_27_mlp_gate.lora_down.weight": [32, 6144], + "lora_unet_blocks_27_mlp_gate.lora_up.weight": [16384, 32], + "lora_unet_blocks_27_mlp_up.alpha": [], + "lora_unet_blocks_27_mlp_up.lora_down.weight": [32, 6144], + "lora_unet_blocks_27_mlp_up.lora_up.weight": [16384, 32], + "lora_unet_first.alpha": [], + "lora_unet_first.lora_down.weight": [32, 64], + "lora_unet_first.lora_up.weight": [6144, 32], + "lora_unet_last_linear.alpha": [], + "lora_unet_last_linear.lora_down.weight": [32, 6144], + "lora_unet_last_linear.lora_up.weight": [64, 32], + "lora_unet_tmlp_0.alpha": [], + "lora_unet_tmlp_0.lora_down.weight": [32, 256], + "lora_unet_tmlp_0.lora_up.weight": [6144, 32], + "lora_unet_tmlp_2.alpha": [], + "lora_unet_tmlp_2.lora_down.weight": [32, 6144], + "lora_unet_tmlp_2.lora_up.weight": [6144, 32], + "lora_unet_tproj_1.alpha": [], + "lora_unet_tproj_1.lora_down.weight": [32, 6144], + "lora_unet_tproj_1.lora_up.weight": [36864, 32], + "lora_unet_txtfusion_layerwise_blocks_0_attn_gate.alpha": [], + "lora_unet_txtfusion_layerwise_blocks_0_attn_gate.lora_down.weight": [32, 2560], + "lora_unet_txtfusion_layerwise_blocks_0_attn_gate.lora_up.weight": [2560, 32], + "lora_unet_txtfusion_layerwise_blocks_0_attn_wk.alpha": [], + "lora_unet_txtfusion_layerwise_blocks_0_attn_wk.lora_down.weight": [32, 2560], + "lora_unet_txtfusion_layerwise_blocks_0_attn_wk.lora_up.weight": [2560, 32], + "lora_unet_txtfusion_layerwise_blocks_0_attn_wo.alpha": [], + "lora_unet_txtfusion_layerwise_blocks_0_attn_wo.lora_down.weight": [32, 2560], + "lora_unet_txtfusion_layerwise_blocks_0_attn_wo.lora_up.weight": [2560, 32], + "lora_unet_txtfusion_layerwise_blocks_0_attn_wq.alpha": [], + "lora_unet_txtfusion_layerwise_blocks_0_attn_wq.lora_down.weight": [32, 2560], + "lora_unet_txtfusion_layerwise_blocks_0_attn_wq.lora_up.weight": [2560, 32], + "lora_unet_txtfusion_layerwise_blocks_0_attn_wv.alpha": [], + "lora_unet_txtfusion_layerwise_blocks_0_attn_wv.lora_down.weight": [32, 2560], + "lora_unet_txtfusion_layerwise_blocks_0_attn_wv.lora_up.weight": [2560, 32], + "lora_unet_txtfusion_layerwise_blocks_0_mlp_down.alpha": [], + "lora_unet_txtfusion_layerwise_blocks_0_mlp_down.lora_down.weight": [32, 6912], + "lora_unet_txtfusion_layerwise_blocks_0_mlp_down.lora_up.weight": [2560, 32], + "lora_unet_txtfusion_layerwise_blocks_0_mlp_gate.alpha": [], + "lora_unet_txtfusion_layerwise_blocks_0_mlp_gate.lora_down.weight": [32, 2560], + "lora_unet_txtfusion_layerwise_blocks_0_mlp_gate.lora_up.weight": [6912, 32], + "lora_unet_txtfusion_layerwise_blocks_0_mlp_up.alpha": [], + "lora_unet_txtfusion_layerwise_blocks_0_mlp_up.lora_down.weight": [32, 2560], + "lora_unet_txtfusion_layerwise_blocks_0_mlp_up.lora_up.weight": [6912, 32], + "lora_unet_txtfusion_projector.alpha": [], + "lora_unet_txtfusion_projector.lora_down.weight": [32, 12], + "lora_unet_txtfusion_projector.lora_up.weight": [1, 32], + "lora_unet_txtfusion_refiner_blocks_1_attn_gate.alpha": [], + "lora_unet_txtfusion_refiner_blocks_1_attn_gate.lora_down.weight": [32, 2560], + "lora_unet_txtfusion_refiner_blocks_1_attn_gate.lora_up.weight": [2560, 32], + "lora_unet_txtfusion_refiner_blocks_1_attn_wk.alpha": [], + "lora_unet_txtfusion_refiner_blocks_1_attn_wk.lora_down.weight": [32, 2560], + "lora_unet_txtfusion_refiner_blocks_1_attn_wk.lora_up.weight": [2560, 32], + "lora_unet_txtfusion_refiner_blocks_1_attn_wo.alpha": [], + "lora_unet_txtfusion_refiner_blocks_1_attn_wo.lora_down.weight": [32, 2560], + "lora_unet_txtfusion_refiner_blocks_1_attn_wo.lora_up.weight": [2560, 32], + "lora_unet_txtfusion_refiner_blocks_1_attn_wq.alpha": [], + "lora_unet_txtfusion_refiner_blocks_1_attn_wq.lora_down.weight": [32, 2560], + "lora_unet_txtfusion_refiner_blocks_1_attn_wq.lora_up.weight": [2560, 32], + "lora_unet_txtfusion_refiner_blocks_1_attn_wv.alpha": [], + "lora_unet_txtfusion_refiner_blocks_1_attn_wv.lora_down.weight": [32, 2560], + "lora_unet_txtfusion_refiner_blocks_1_attn_wv.lora_up.weight": [2560, 32], + "lora_unet_txtfusion_refiner_blocks_1_mlp_down.alpha": [], + "lora_unet_txtfusion_refiner_blocks_1_mlp_down.lora_down.weight": [32, 6912], + "lora_unet_txtfusion_refiner_blocks_1_mlp_down.lora_up.weight": [2560, 32], + "lora_unet_txtfusion_refiner_blocks_1_mlp_gate.alpha": [], + "lora_unet_txtfusion_refiner_blocks_1_mlp_gate.lora_down.weight": [32, 2560], + "lora_unet_txtfusion_refiner_blocks_1_mlp_gate.lora_up.weight": [6912, 32], + "lora_unet_txtfusion_refiner_blocks_1_mlp_up.alpha": [], + "lora_unet_txtfusion_refiner_blocks_1_mlp_up.lora_down.weight": [32, 2560], + "lora_unet_txtfusion_refiner_blocks_1_mlp_up.lora_up.weight": [6912, 32], + "lora_unet_txtmlp_1.alpha": [], + "lora_unet_txtmlp_1.lora_down.weight": [32, 2560], + "lora_unet_txtmlp_1.lora_up.weight": [6144, 32], + "lora_unet_txtmlp_3.alpha": [], + "lora_unet_txtmlp_3.lora_down.weight": [32, 6144], + "lora_unet_txtmlp_3.lora_up.weight": [6144, 32], +} 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..10c2632987b 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 @@ -1,6 +1,9 @@ +import accelerate import pytest import torch +from diffusers import Krea2Transformer2DModel +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.lora_layer import LoRALayer from invokeai.backend.patches.lora_conversions.krea2_lora_constants import ( @@ -8,6 +11,10 @@ KREA2_LORA_TRANSFORMER_PREFIX, ) 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, +) +from tests.backend.patches.lora_conversions.lora_state_dicts.utils import keys_to_mock_state_dict def test_peft_layer_preserves_explicit_alpha() -> None: @@ -208,6 +215,182 @@ def test_single_module_native_krea2_lora_is_remapped( assert set(model.layers) == {f"{KREA2_LORA_TRANSFORMER_PREFIX}{diffusers_module}"} +@pytest.mark.parametrize( + ("kohya_module", "diffusers_module"), + [ + ("blocks_0_attn_wq", "transformer_blocks.0.attn.to_q"), + ("blocks_0_attn_wk", "transformer_blocks.0.attn.to_k"), + ("blocks_0_attn_wv", "transformer_blocks.0.attn.to_v"), + ("blocks_0_attn_wo", "transformer_blocks.0.attn.to_out.0"), + ("blocks_0_attn_gate", "transformer_blocks.0.attn.to_gate"), + # Multi-digit block index. + ("blocks_27_mlp_down", "transformer_blocks.27.ff.down"), + # `layerwise_blocks` / `refiner_blocks` are the native components that themselves contain an + # underscore, i.e. the only genuine ambiguity in the flattened form. + ("txtfusion_layerwise_blocks_0_attn_wo", "text_fusion.layerwise_blocks.0.attn.to_out.0"), + ("txtfusion_refiner_blocks_1_mlp_gate", "text_fusion.refiner_blocks.1.ff.gate"), + ("txtfusion_projector", "text_fusion.projector"), + ("first", "img_in"), + ("tmlp_0", "time_embed.linear_1"), + ("tmlp_2", "time_embed.linear_2"), + ("tproj_1", "time_mod_proj"), + ("txtmlp_1", "txt_in.linear_1"), + ("txtmlp_3", "txt_in.linear_2"), + ("last_linear", "final_layer.linear"), + ], +) +def test_kohya_flattened_krea2_module_is_remapped(kohya_module: str, diffusers_module: str) -> None: + # kohya / LyCORIS flatten the module path and prefix it with `lora_unet_`. Without un-flattening, every key + # misses its module and the adapter is a silent no-op ("Failed to find module for LoRA layer key: + # lora_transformer-lora_unet_blocks_6_attn_wv"). + state_dict = { + f"lora_unet_{kohya_module}.lora_down.weight": torch.ones(2, 4), + f"lora_unet_{kohya_module}.lora_up.weight": torch.ones(4, 2), + } + + model = lora_model_from_krea2_state_dict(state_dict) + + assert set(model.layers) == {f"{KREA2_LORA_TRANSFORMER_PREFIX}{diffusers_module}"} + + +def test_kohya_krea2_lora_layers_match_the_real_transformer() -> None: + # Every layer of a real kohya Krea-2 adapter must land on an actual Linear of Krea2Transformer2DModel, with + # in/out features that agree with the LoRA's own down/up shapes. A wrong rename (e.g. ff.gate <-> ff.down, + # whose SwiGLU shapes are transposed) is caught here rather than as a runtime warning. + state_dict = keys_to_mock_state_dict(krea2_kohya_state_dict_keys) + + model = lora_model_from_krea2_state_dict(state_dict) + + with accelerate.init_empty_weights(): + transformer = Krea2Transformer2DModel(**KREA2_TRANSFORMER_CONFIG) + + # 4 blocks x 8 Linears (2 transformer, 1 layerwise + 1 refiner text-fusion) + 8 top-level modules. + assert len(model.layers) == 40 + for layer_key, layer in model.layers.items(): + module_name = layer_key[len(KREA2_LORA_TRANSFORMER_PREFIX) :] + submodule = transformer.get_submodule(module_name) + assert isinstance(submodule, torch.nn.Linear), f"{module_name} is not a Linear" + out_features, in_features = submodule.weight.shape + assert layer.down.shape[1] == in_features, f"{module_name} in_features mismatch" + assert layer.up.shape[0] == out_features, f"{module_name} out_features mismatch" + + +def test_kohya_flattened_krea2_layer_preserves_alpha() -> None: + # kohya adapters carry an explicit `.alpha`; it must survive the un-flattening intact, otherwise the LoRA + # applies at the wrong strength. + state_dict = { + "lora_unet_blocks_6_attn_wv.lora_down.weight": torch.ones(2, 4), + "lora_unet_blocks_6_attn_wv.lora_up.weight": torch.ones(4, 2), + "lora_unet_blocks_6_attn_wv.alpha": torch.tensor(2.0), + } + + model = lora_model_from_krea2_state_dict(state_dict) + + layer = model.layers[f"{KREA2_LORA_TRANSFORMER_PREFIX}transformer_blocks.6.attn.to_v"] + assert isinstance(layer, LoRALayer) + assert layer._alpha == 2.0 + + +def test_kohya_flattened_krea2_keys_tolerate_doubled_separator() -> None: + state_dict = { + "lora_unet__blocks_6_attn_wv.lora_down.weight": torch.ones(2, 4), + "lora_unet__blocks_6_attn_wv.lora_up.weight": torch.ones(4, 2), + } + + model = lora_model_from_krea2_state_dict(state_dict) + + assert set(model.layers) == {f"{KREA2_LORA_TRANSFORMER_PREFIX}transformer_blocks.6.attn.to_v"} + + +@pytest.mark.parametrize( + "flat_module", + [ + # Non-Linear natives: `mod.lin` is folded into the `scale_shift_table` parameter and the norms have no + # Linear counterpart, so there is nothing to patch. They must not be renamed into a key that merely + # looks applicable. + "blocks_6_mod_lin", + "blocks_6_prenorm", + "blocks_6_attn_qknorm_qnorm", + # Not a Krea-2 module layout at all (e.g. a flattened adapter for some other architecture). + "double_blocks_0_img_attn_proj", + # Sequential positions that hold an activation rather than a Linear. The parsing tree enumerates the + # indices it accepts, so these are rejected outright instead of being rewritten to `tmlp.1.*` — a + # half-converted key that the native pass no longer recognizes. + "tmlp_1", + "tproj_0", + "txtmlp_2", + ], +) +def test_unrecognized_kohya_flattened_keys_are_left_untouched(flat_module: str) -> None: + state_dict = { + f"lora_unet_{flat_module}.lora_down.weight": torch.ones(2, 4), + f"lora_unet_{flat_module}.lora_up.weight": torch.ones(4, 2), + } + + model = lora_model_from_krea2_state_dict(state_dict) + + # Left verbatim rather than rewritten into a plausible-looking but wrong module path. + assert set(model.layers) == {f"{KREA2_LORA_TRANSFORMER_PREFIX}lora_unet_{flat_module}"} + + +@pytest.mark.parametrize( + "lycoris_suffixes", + [ + ("lokr_w1", "lokr_w2"), + ("hada_w1_a", "hada_w1_b", "hada_w2_a", "hada_w2_b"), + ("diff", "diff_b"), + ], +) +def test_kohya_lycoris_algorithm_keys_do_not_abort_the_load(lycoris_suffixes: tuple[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 + # two modules into one unsupported layer — aborting the *entire* adapter at generation time. + state_dict = { + "lora_unet_blocks_0_attn_wv.lora_down.weight": torch.ones(2, 4), + "lora_unet_blocks_0_attn_wv.lora_up.weight": torch.ones(4, 2), + **{f"lora_unet_blocks_6_attn_wq.{suffix}": torch.ones(4, 4) for suffix in lycoris_suffixes}, + } + + 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. + 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", + } + + +def test_non_string_keys_survive_the_kohya_and_native_passes() -> None: + # `.pt`/`.ckpt` sources can carry non-string keys. Once the kohya pass rewrites something, the native pass + # runs its substring tests over every key — which raised `TypeError: argument of type 'int' is not + # iterable` on an int key rather than leaving it alone. + state_dict = { + 0: torch.ones(2), + "lora_unet_blocks_0_attn_wv.lora_down.weight": torch.ones(2, 4), + "lora_unet_blocks_0_attn_wv.lora_up.weight": torch.ones(4, 2), + } + + model = lora_model_from_krea2_state_dict(state_dict) + + assert f"{KREA2_LORA_TRANSFORMER_PREFIX}transformer_blocks.0.attn.to_v" in model.layers + + +def test_conflicting_kohya_and_native_aliases_raise() -> None: + # The flattened and dotted spellings of one logical layer normalize to the same target. Providing both + # must raise instead of silently dropping one based on dict ordering. + state_dict = { + "lora_unet_blocks_0_attn_wq.lora_down.weight": torch.ones(2, 4), + "lora_unet_blocks_0_attn_wq.lora_up.weight": torch.ones(4, 2), + "blocks.0.attn.wq.lora_down.weight": torch.full((2, 4), 2.0), + "blocks.0.attn.wq.lora_up.weight": torch.full((4, 2), 2.0), + } + + with pytest.raises(ValueError, match="normalize to the same target"): + lora_model_from_krea2_state_dict(state_dict) + + def test_native_transformer_remap_does_not_change_diffusers_text_encoder_blocks() -> None: state_dict = { "diffusion_model.blocks.0.attn.wq.lora_A.weight": torch.ones(2, 4),