Skip to content
Open
Show file tree
Hide file tree
Changes from 9 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions invokeai/backend/model_manager/configs/lora.py
Original file line number Diff line number Diff line change
Expand Up @@ -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_<path>` (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.",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,9 @@
``transformer.<module>.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_<path>``.
"""

import re
Expand All @@ -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,
Expand Down Expand Up @@ -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}'. "
Expand All @@ -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.

Expand All @@ -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)

Expand Down
36 changes: 36 additions & 0 deletions tests/backend/model_manager/configs/test_krea2_lora_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
Original file line number Diff line number Diff line change
@@ -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],
}
Loading
Loading