diff --git a/invokeai/backend/model_manager/configs/lora.py b/invokeai/backend/model_manager/configs/lora.py index bf51a9734e6..149e27ae16e 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, @@ -996,6 +996,56 @@ def _has_complete_lora_pair(state_dict: dict[str | int, Any], key_filter: Callab return False +# 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 _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) or not key.endswith(_LOKR_WEIGHT_SUFFIXES): + continue + if key_filter is not None and not key_filter(key): + continue + 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 # 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`. @@ -1044,10 +1094,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(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): + has_supported_explicit_layer = _has_complete_lora_pair( + state_dict, _key_is_supported_krea2_layout + ) 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) @@ -1057,18 +1115,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. - if not (_has_krea2_lora_keys(state_dict) and _has_complete_lora_pair(state_dict)): + # 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_complete_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. + # 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 2b5f66a8093..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: @@ -237,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( @@ -256,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 @@ -267,6 +285,9 @@ def _get_lora_layer_values( # dim -> value key ``dora_magnitude`` # Mapping them here lets a 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", @@ -276,6 +297,7 @@ def _get_lora_layer_values( ".lora_magnitude_vector.weight": "dora_magnitude", ".magnitude": "dora_magnitude", ".alpha": "alpha", + **{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 b10064d887e..fe7fa052312 100644 --- a/tests/backend/model_manager/configs/test_krea2_lora_config.py +++ b/tests/backend/model_manager/configs/test_krea2_lora_config.py @@ -371,3 +371,118 @@ 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 + + +_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. + 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") +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 + + +@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_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( + _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 = _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 fbbb745fa07..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 @@ -5,6 +5,7 @@ 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, @@ -370,25 +371,31 @@ 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"), - # `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. + (("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"), + # 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(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 @@ -404,11 +411,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}", } @@ -483,3 +491,104 @@ 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_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 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": w1, + "transformer.text_fusion.0.attn.to_q.lokr_w2": w2, + "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) + 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_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": 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) + 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": 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) + + 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)) + + +@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 = { + "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), + } + + 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) + 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": 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), + } + + 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) + 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]