diff --git a/invokeai/backend/patches/lora_conversions/anima_lora_conversion_utils.py b/invokeai/backend/patches/lora_conversions/anima_lora_conversion_utils.py index b55a96dca75..3258b2da4bd 100644 --- a/invokeai/backend/patches/lora_conversions/anima_lora_conversion_utils.py +++ b/invokeai/backend/patches/lora_conversions/anima_lora_conversion_utils.py @@ -140,10 +140,10 @@ def _make_layer_patch(layer_dict: dict[str, torch.Tensor]) -> BaseLayerPatch: from LoKR layers so they fall through to the LoKR handler instead. """ has_lokr = "lokr_w1" in layer_dict or "lokr_w1_a" in layer_dict - has_dora = "dora_scale" in layer_dict - if has_lokr and has_dora: - layer_dict = {k: v for k, v in layer_dict.items() if k != "dora_scale"} - logger.warning("Stripped dora_scale from LoKR layer (DoRA+LoKR combination not supported, using LoKR only)") + dora_keys = {"dora_scale", "dora_magnitude"}.intersection(layer_dict) + if has_lokr and dora_keys: + layer_dict = {k: v for k, v in layer_dict.items() if k not in dora_keys} + logger.warning("Stripped DoRA magnitude from LoKR layer (DoRA+LoKR combination not supported, using LoKR only)") return any_lora_layer_from_state_dict(layer_dict) @@ -154,6 +154,8 @@ def _make_layer_patch(layer_dict: dict[str, torch.Tensor]) -> BaseLayerPatch: ".lora_down.weight", ".lora_up.weight", ".dora_scale", + ".lora_magnitude_vector.weight", + ".magnitude", ".alpha", ] @@ -197,6 +199,8 @@ def _group_keys_by_layer( if key.endswith(suffix): layer_name = key[: -len(suffix)] key_name = suffix[1:] # Remove leading dot + if key_name in {"lora_magnitude_vector.weight", "magnitude"}: + key_name = "dora_magnitude" break if layer_name is None: @@ -220,6 +224,9 @@ def _get_lora_layer_values(layer_dict: dict[str, torch.Tensor], alpha: float | N } if alpha is not None: values["alpha"] = torch.tensor(alpha) + for magnitude_key in ("dora_scale", "dora_magnitude"): + if magnitude_key in layer_dict: + values[magnitude_key] = layer_dict[magnitude_key] return values elif "lora_down.weight" in layer_dict: return layer_dict diff --git a/invokeai/backend/patches/lora_conversions/flux_aitoolkit_lora_conversion_utils.py b/invokeai/backend/patches/lora_conversions/flux_aitoolkit_lora_conversion_utils.py index f359e7caa32..912bde71abd 100644 --- a/invokeai/backend/patches/lora_conversions/flux_aitoolkit_lora_conversion_utils.py +++ b/invokeai/backend/patches/lora_conversions/flux_aitoolkit_lora_conversion_utils.py @@ -30,8 +30,9 @@ def is_state_dict_likely_in_flux_aitoolkit_format( if not _has_flux_layer_structure(state_dict): return False - # AIToolkit only produces standard PEFT LoRA (lora_A.weight / lora_B.weight). - # Exclude LyCORIS algorithm variants (LoKR, LoHA, etc.) which use different weight key suffixes. + # AIToolkit produces standard PEFT LoRA (lora_A.weight / lora_B.weight) and DoRA + # magnitude sidecars. Exclude LyCORIS algorithm variants (LoKR, LoHA, etc.) which + # use different weight key suffixes. # These are handled by the BFL PEFT converter instead. _LYCORIS_SUFFIXES = ( "lokr_w1", diff --git a/invokeai/backend/patches/lora_conversions/flux_bfl_peft_lora_conversion_utils.py b/invokeai/backend/patches/lora_conversions/flux_bfl_peft_lora_conversion_utils.py index fd89d673c8f..b14b9a3c72b 100644 --- a/invokeai/backend/patches/lora_conversions/flux_bfl_peft_lora_conversion_utils.py +++ b/invokeai/backend/patches/lora_conversions/flux_bfl_peft_lora_conversion_utils.py @@ -53,7 +53,7 @@ _SINGLE_BLOCK_RE = re.compile(r"^single_blocks\.(\d+)\.(.+)$") # Weight key suffixes used by PEFT LoRA in BFL format. -_BFL_PEFT_LORA_SUFFIXES = ("lora_A.weight", "lora_B.weight") +_BFL_PEFT_LORA_SUFFIXES = ("lora_A.weight", "lora_B.weight", "lora_magnitude_vector.weight", "magnitude") # Weight key suffixes used by LyCORIS algorithms (LoKR, LoHA, etc.) in BFL format. # These are single-component suffixes (no dot), unlike the two-component PEFT suffixes. @@ -193,6 +193,8 @@ def lora_model_from_flux_bfl_peft_state_dict( grouped_state_dict[layer_name]["lora_down.weight"] = value elif suffix == "lora_B.weight": grouped_state_dict[layer_name]["lora_up.weight"] = value + elif suffix in ("lora_magnitude_vector.weight", "magnitude"): + grouped_state_dict[layer_name]["dora_magnitude"] = value else: grouped_state_dict[layer_name][suffix] = value @@ -238,6 +240,8 @@ def lora_model_from_flux2_bfl_peft_state_dict( grouped_state_dict[layer_name]["lora_down.weight"] = value elif suffix == "lora_B.weight": grouped_state_dict[layer_name]["lora_up.weight"] = value + elif suffix in ("lora_magnitude_vector.weight", "magnitude"): + grouped_state_dict[layer_name]["dora_magnitude"] = value else: grouped_state_dict[layer_name][suffix] = value @@ -347,18 +351,25 @@ def _split_qkv_lora( lora_down = layer_sd["lora_down.weight"] # [rank, hidden] lora_up = layer_sd["lora_up.weight"] # [3*hidden, rank] alpha = layer_sd.get("alpha") + dora_scale = layer_sd.get("dora_scale") + dora_magnitude = layer_sd.get("dora_magnitude") # Split lora_up into 3 equal parts along dim 0 up_q, up_k, up_v = lora_up.chunk(3, dim=0) + magnitude_parts = dora_magnitude.chunk(3, dim=0) if dora_magnitude is not None else (None, None, None) result = [] - for key, up_part in [(q_key, up_q), (k_key, up_k), (v_key, up_v)]: + for key, up_part, magnitude_part in zip((q_key, k_key, v_key), (up_q, up_k, up_v), magnitude_parts, strict=True): sd: dict[str, torch.Tensor] = { "lora_down.weight": lora_down.clone(), "lora_up.weight": up_part, } if alpha is not None: sd["alpha"] = alpha + if dora_scale is not None: + sd["dora_scale"] = dora_scale.clone() + if magnitude_part is not None: + sd["dora_magnitude"] = magnitude_part result.append((key, sd)) return result diff --git a/invokeai/backend/patches/lora_conversions/flux_diffusers_lora_conversion_utils.py b/invokeai/backend/patches/lora_conversions/flux_diffusers_lora_conversion_utils.py index 05fe4cab297..25526491e3f 100644 --- a/invokeai/backend/patches/lora_conversions/flux_diffusers_lora_conversion_utils.py +++ b/invokeai/backend/patches/lora_conversions/flux_diffusers_lora_conversion_utils.py @@ -8,6 +8,17 @@ from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_TRANSFORMER_PREFIX from invokeai.backend.patches.model_patch_raw import ModelPatchRaw +_LORA_SUFFIX_TO_VALUE_KEY = { + ".lora_A.weight": "lora_A.weight", + ".lora_B.weight": "lora_B.weight", + ".lora.down.weight": "lora.down.weight", + ".lora.up.weight": "lora.up.weight", + ".dora_scale": "dora_scale", + ".lora_magnitude_vector.weight": "dora_magnitude", + ".magnitude": "dora_magnitude", + ".alpha": "alpha", +} + def is_state_dict_likely_in_flux_diffusers_format(state_dict: dict[str | int, torch.Tensor]) -> bool: """Checks if the provided state dict is likely in the Diffusers FLUX LoRA format. @@ -20,7 +31,14 @@ def is_state_dict_likely_in_flux_diffusers_format(state_dict: dict[str | int, to """ # Check that all keys are LoRA weight keys (either PEFT or standard format). # Some LoRAs use a mix of formats (PEFT for some layers, standard for others). - _LORA_SUFFIXES = ("lora_A.weight", "lora_B.weight", "lora.down.weight", "lora.up.weight") + _LORA_SUFFIXES = ( + "lora_A.weight", + "lora_B.weight", + "lora.down.weight", + "lora.up.weight", + "lora_magnitude_vector.weight", + "magnitude", + ) all_keys_are_lora = all(k.endswith(_LORA_SUFFIXES) for k in state_dict.keys() if isinstance(k, str)) if not all_keys_are_lora: return False @@ -125,12 +143,22 @@ def get_lora_layer_values(src_layer_dict: dict[str, torch.Tensor]) -> dict[str, if "lora_A.weight" in src_layer_dict: # The LoRA keys are in PEFT format. values = { - "lora_down.weight": src_layer_dict.pop("lora_A.weight"), - "lora_up.weight": src_layer_dict.pop("lora_B.weight"), + "lora_down.weight": src_layer_dict["lora_A.weight"], + "lora_up.weight": src_layer_dict["lora_B.weight"], } if alpha is not None: values["alpha"] = torch.tensor(alpha) - assert len(src_layer_dict) == 0 + for value_key in ("alpha", "dora_scale", "dora_magnitude"): + if value_key in src_layer_dict and (value_key != "alpha" or alpha is None): + values[value_key] = src_layer_dict[value_key] + unexpected_keys = set(src_layer_dict) - { + "lora_A.weight", + "lora_B.weight", + "alpha", + "dora_scale", + "dora_magnitude", + } + assert not unexpected_keys, f"Unexpected LoRA keys: {unexpected_keys}" return values else: # Assume that the LoRA keys are in Kohya format. @@ -327,6 +355,13 @@ def lora_model_from_flux2_diffusers_state_dict( else: values = src_layer_dict + for source_key in ("lora_magnitude_vector.weight", "magnitude"): + if source_key in src_layer_dict: + values["dora_magnitude"] = src_layer_dict[source_key] + for magnitude_key in ("dora_scale", "dora_magnitude"): + if magnitude_key in src_layer_dict: + values[magnitude_key] = src_layer_dict[magnitude_key] + if alpha is not None and "alpha" not in values: values["alpha"] = torch.tensor(alpha) @@ -347,11 +382,19 @@ def _group_by_layer_mixed_format(state_dict: Dict[str, torch.Tensor]) -> dict[st continue # Determine suffix length based on the key ending - if key.endswith((".lora_A.weight", ".lora_B.weight")): + if key.endswith((".lora_A.weight", ".lora_B.weight", ".lora_magnitude_vector.weight")): # PEFT format: split off 2 parts (lora_A + weight) parts = key.rsplit(".", maxsplit=2) layer_name = parts[0] suffix = ".".join(parts[1:]) + if suffix == "lora_magnitude_vector.weight": + suffix = "dora_magnitude" + elif key.endswith(".magnitude"): + layer_name = key[: -len(".magnitude")] + suffix = "dora_magnitude" + elif key.endswith(".dora_scale"): + layer_name = key[: -len(".dora_scale")] + suffix = "dora_scale" elif key.endswith((".lora.down.weight", ".lora.up.weight")): # Standard format: split off 3 parts (lora + down/up + weight) parts = key.rsplit(".", maxsplit=3) @@ -374,10 +417,17 @@ def _group_by_layer(state_dict: Dict[str, torch.Tensor]) -> dict[str, dict[str, """Groups the keys in the state dict by layer.""" layer_dict: dict[str, dict[str, torch.Tensor]] = {} for key in state_dict: - # Split the 'lora_A.weight' or 'lora_B.weight' suffix from the layer name. - parts = key.rsplit(".", maxsplit=2) - layer_name = parts[0] - key_name = ".".join(parts[1:]) + layer_name = None + key_name = None + for suffix, value_key in _LORA_SUFFIX_TO_VALUE_KEY.items(): + if key.endswith(suffix): + layer_name = key[: -len(suffix)] + key_name = value_key + break + if layer_name is None: + parts = key.rsplit(".", maxsplit=2) + layer_name = parts[0] + key_name = ".".join(parts[1:]) if layer_name not in layer_dict: layer_dict[layer_name] = {} layer_dict[layer_name][key_name] = state_dict[key] diff --git a/invokeai/backend/patches/lora_conversions/flux_kohya_lora_conversion_utils.py b/invokeai/backend/patches/lora_conversions/flux_kohya_lora_conversion_utils.py index f5a6830c4f1..55a4e0a4178 100644 --- a/invokeai/backend/patches/lora_conversions/flux_kohya_lora_conversion_utils.py +++ b/invokeai/backend/patches/lora_conversions/flux_kohya_lora_conversion_utils.py @@ -99,6 +99,9 @@ def lora_model_from_flux_kohya_state_dict(state_dict: Dict[str, torch.Tensor]) - (FLUX_LORA_T5_PREFIX, t5_grouped_sd), ]: for layer_key, layer_state_dict in grouped_sd.items(): + for source_key in ("magnitude", "lora_magnitude_vector.weight"): + if source_key in layer_state_dict: + layer_state_dict["dora_magnitude"] = layer_state_dict.pop(source_key) layers[model_prefix + layer_key] = any_lora_layer_from_state_dict(layer_state_dict) # Create and return the LoRAModelRaw. diff --git a/invokeai/backend/patches/lora_conversions/flux_onetrainer_bfl_lora_conversion_utils.py b/invokeai/backend/patches/lora_conversions/flux_onetrainer_bfl_lora_conversion_utils.py index b2109222a31..fbe054795f5 100644 --- a/invokeai/backend/patches/lora_conversions/flux_onetrainer_bfl_lora_conversion_utils.py +++ b/invokeai/backend/patches/lora_conversions/flux_onetrainer_bfl_lora_conversion_utils.py @@ -32,7 +32,13 @@ _TRANSFORMER_PREFIX = "transformer." # Valid LoRA weight suffixes in this format. -_LORA_SUFFIXES = ("lora_down.weight", "lora_up.weight", "alpha") +_LORA_SUFFIXES = ( + "lora_down.weight", + "lora_up.weight", + "alpha", + "lora_magnitude_vector.weight", + "magnitude", +) # Regex to detect split QKV keys in double blocks: e.g. "double_blocks.0.img_attn.qkv.1" _SPLIT_QKV_RE = re.compile(r"^(double_blocks\.\d+\.(img_attn|txt_attn)\.qkv)\.\d+$") @@ -111,6 +117,8 @@ def lora_model_from_flux_onetrainer_bfl_state_dict(state_dict: Dict[str, torch.T if layer_name not in grouped_state_dict: grouped_state_dict[layer_name] = {} + if suffix in ("magnitude", "lora_magnitude_vector.weight"): + suffix = "dora_magnitude" grouped_state_dict[layer_name][suffix] = value # Step 2: Build LoRA layers, merging split QKV and linear1. diff --git a/invokeai/backend/patches/lora_conversions/flux_onetrainer_lora_conversion_utils.py b/invokeai/backend/patches/lora_conversions/flux_onetrainer_lora_conversion_utils.py index 88aeee95e49..5b9545f2b05 100644 --- a/invokeai/backend/patches/lora_conversions/flux_onetrainer_lora_conversion_utils.py +++ b/invokeai/backend/patches/lora_conversions/flux_onetrainer_lora_conversion_utils.py @@ -92,6 +92,9 @@ def lora_model_from_flux_onetrainer_state_dict(state_dict: Dict[str, torch.Tenso (FLUX_LORA_T5_PREFIX, t5_grouped_sd), ]: for layer_key, layer_state_dict in grouped_sd.items(): + for source_key in ("magnitude", "lora_magnitude_vector.weight"): + if source_key in layer_state_dict: + layer_state_dict["dora_magnitude"] = layer_state_dict.pop(source_key) layers[model_prefix + layer_key] = any_lora_layer_from_state_dict(layer_state_dict) # Handle the transformer. diff --git a/invokeai/backend/patches/lora_conversions/qwen_image_lora_conversion_utils.py b/invokeai/backend/patches/lora_conversions/qwen_image_lora_conversion_utils.py index 7fc01f72315..04d748c17b0 100644 --- a/invokeai/backend/patches/lora_conversions/qwen_image_lora_conversion_utils.py +++ b/invokeai/backend/patches/lora_conversions/qwen_image_lora_conversion_utils.py @@ -114,7 +114,8 @@ def _convert_kohya_format(state_dict: Dict[str, torch.Tensor], alpha: float | No if model_path is None: continue # Skip unrecognized layers - layer = any_lora_layer_from_state_dict(layer_dict) + values = _normalize_lora_keys(layer_dict, alpha) + layer = any_lora_layer_from_state_dict(values) final_key = f"{QWEN_IMAGE_EDIT_LORA_TRANSFORMER_PREFIX}{model_path}" layers[final_key] = layer @@ -146,6 +147,11 @@ def _convert_diffusers_format(state_dict: Dict[str, torch.Tensor], alpha: float def _normalize_lora_keys(layer_dict: dict[str, torch.Tensor], alpha: float | None) -> dict[str, torch.Tensor]: """Normalize LoRA key names to internal format.""" + layer_dict = dict(layer_dict) + for source_key in ("lora_magnitude_vector.weight", "magnitude"): + if source_key in layer_dict: + layer_dict["dora_magnitude"] = layer_dict.pop(source_key) + if "lora_A.weight" in layer_dict: values: dict[str, torch.Tensor] = { "lora_down.weight": layer_dict["lora_A.weight"], @@ -153,6 +159,9 @@ def _normalize_lora_keys(layer_dict: dict[str, torch.Tensor], alpha: float | Non } if alpha is not None: values["alpha"] = torch.tensor(alpha) + for magnitude_key in ("dora_scale", "dora_magnitude"): + if magnitude_key in layer_dict: + values[magnitude_key] = layer_dict[magnitude_key] return values elif "lora_down.weight" in layer_dict: return layer_dict @@ -164,14 +173,16 @@ def _group_by_layer(state_dict: Dict[str, torch.Tensor]) -> dict[str, dict[str, """Group state dict keys by layer path.""" layer_dict: dict[str, dict[str, torch.Tensor]] = {} - known_suffixes = [ - ".lora_A.weight", - ".lora_B.weight", - ".lora_down.weight", - ".lora_up.weight", - ".dora_scale", - ".alpha", - ] + suffix_to_value_key = { + ".lora_A.weight": "lora_A.weight", + ".lora_B.weight": "lora_B.weight", + ".lora_down.weight": "lora_down.weight", + ".lora_up.weight": "lora_up.weight", + ".dora_scale": "dora_scale", + ".lora_magnitude_vector.weight": "dora_magnitude", + ".magnitude": "dora_magnitude", + ".alpha": "alpha", + } for key in state_dict: if not isinstance(key, str): @@ -179,10 +190,10 @@ def _group_by_layer(state_dict: Dict[str, torch.Tensor]) -> dict[str, dict[str, layer_name = None key_name = None - for suffix in known_suffixes: + for suffix, value_key in suffix_to_value_key.items(): if key.endswith(suffix): layer_name = key[: -len(suffix)] - key_name = suffix[1:] + key_name = value_key break if layer_name is None: diff --git a/invokeai/backend/patches/lora_conversions/wan_lora_conversion_utils.py b/invokeai/backend/patches/lora_conversions/wan_lora_conversion_utils.py index 5592572b246..4f1976c3639 100644 --- a/invokeai/backend/patches/lora_conversions/wan_lora_conversion_utils.py +++ b/invokeai/backend/patches/lora_conversions/wan_lora_conversion_utils.py @@ -211,6 +211,8 @@ def _normalize_lora_param_names(layer_dict: dict[str, torch.Tensor], alpha: floa values["alpha"] = layer_dict["alpha"] if "dora_scale" in layer_dict: values["dora_scale"] = layer_dict["dora_scale"] + if "dora_magnitude" in layer_dict: + values["dora_magnitude"] = layer_dict["dora_magnitude"] return values return layer_dict @@ -219,14 +221,16 @@ def _group_by_layer(state_dict: Dict[str, torch.Tensor]) -> dict[str, dict[str, """Group state-dict keys by their layer path (everything before the LoRA-suffix tail).""" grouped: dict[str, dict[str, torch.Tensor]] = {} - known_suffixes = [ - ".lora_A.weight", - ".lora_B.weight", - ".lora_down.weight", - ".lora_up.weight", - ".dora_scale", - ".alpha", - ] + suffix_to_value_key = { + ".lora_A.weight": "lora_A.weight", + ".lora_B.weight": "lora_B.weight", + ".lora_down.weight": "lora_down.weight", + ".lora_up.weight": "lora_up.weight", + ".dora_scale": "dora_scale", + ".lora_magnitude_vector.weight": "dora_magnitude", + ".magnitude": "dora_magnitude", + ".alpha": "alpha", + } for key in state_dict: if not isinstance(key, str): @@ -234,10 +238,10 @@ def _group_by_layer(state_dict: Dict[str, torch.Tensor]) -> dict[str, dict[str, layer_name = None key_name = None - for suffix in known_suffixes: + for suffix, value_key in suffix_to_value_key.items(): if key.endswith(suffix): layer_name = key[: -len(suffix)] - key_name = suffix[1:] # drop leading dot + key_name = value_key break if layer_name is None: diff --git a/invokeai/backend/patches/lora_conversions/z_image_lora_conversion_utils.py b/invokeai/backend/patches/lora_conversions/z_image_lora_conversion_utils.py index 70b10de50d6..4b4d2cd4123 100644 --- a/invokeai/backend/patches/lora_conversions/z_image_lora_conversion_utils.py +++ b/invokeai/backend/patches/lora_conversions/z_image_lora_conversion_utils.py @@ -201,6 +201,9 @@ def _get_lora_layer_values(layer_dict: dict[str, torch.Tensor], alpha: float | N } if alpha is not None: values["alpha"] = torch.tensor(alpha) + for magnitude_key in ("dora_scale", "dora_magnitude"): + if magnitude_key in layer_dict: + values[magnitude_key] = layer_dict[magnitude_key] return values elif "lora_down.weight" in layer_dict: # Already in internal format @@ -225,14 +228,16 @@ def _group_by_layer(state_dict: Dict[str, torch.Tensor]) -> dict[str, dict[str, layer_dict: dict[str, dict[str, torch.Tensor]] = {} # Known suffixes that indicate the end of a layer name - known_suffixes = [ - ".lora_A.weight", - ".lora_B.weight", - ".lora_down.weight", - ".lora_up.weight", - ".dora_scale", - ".alpha", - ] + suffix_to_value_key = { + ".lora_A.weight": "lora_A.weight", + ".lora_B.weight": "lora_B.weight", + ".lora_down.weight": "lora_down.weight", + ".lora_up.weight": "lora_up.weight", + ".dora_scale": "dora_scale", + ".lora_magnitude_vector.weight": "dora_magnitude", + ".magnitude": "dora_magnitude", + ".alpha": "alpha", + } for key in state_dict: if not isinstance(key, str): @@ -241,10 +246,10 @@ def _group_by_layer(state_dict: Dict[str, torch.Tensor]) -> dict[str, dict[str, # Try to find a known suffix layer_name = None key_name = None - for suffix in known_suffixes: + for suffix, value_key in suffix_to_value_key.items(): if key.endswith(suffix): layer_name = key[: -len(suffix)] - key_name = suffix[1:] # Remove leading dot + key_name = value_key break if layer_name is None: diff --git a/tests/backend/patches/lora_conversions/test_anima_lora_conversion_utils.py b/tests/backend/patches/lora_conversions/test_anima_lora_conversion_utils.py index 3d7972cda38..0905ea971ea 100644 --- a/tests/backend/patches/lora_conversions/test_anima_lora_conversion_utils.py +++ b/tests/backend/patches/lora_conversions/test_anima_lora_conversion_utils.py @@ -1,6 +1,8 @@ import pytest import torch +from invokeai.backend.patches.layers.dora_layer import DoRALayer +from invokeai.backend.patches.layers.lokr_layer import LoKRLayer from invokeai.backend.patches.lora_conversions.anima_lora_constants import ( ANIMA_LORA_QWEN3_PREFIX, ANIMA_LORA_TRANSFORMER_PREFIX, @@ -228,3 +230,37 @@ def test_empty_state_dict_returns_empty_model(): """An empty state dict should produce a ModelPatchRaw with no layers.""" lora_model = lora_model_from_anima_state_dict({}) assert len(lora_model.layers) == 0 + + +@pytest.mark.parametrize("magnitude_suffix", ["magnitude", "lora_magnitude_vector.weight"]) +def test_peft_dora_magnitude_is_preserved(magnitude_suffix: str): + in_dim, out_dim, rank = 5, 7, 2 + magnitude = torch.arange(1, out_dim + 1, dtype=torch.float32) + prefix = "diffusion_model.blocks.0.cross_attn.q_proj" + state_dict = { + f"{prefix}.lora_A.weight": torch.zeros(rank, in_dim), + f"{prefix}.lora_B.weight": torch.zeros(out_dim, rank), + f"{prefix}.{magnitude_suffix}": magnitude, + } + + model = lora_model_from_anima_state_dict(state_dict) + layer = model.layers[f"{ANIMA_LORA_TRANSFORMER_PREFIX}blocks.0.cross_attn.q_proj"] + + assert isinstance(layer, DoRALayer) + assert layer.magnitude_is_out_dim is True + assert torch.equal(layer.dora_scale, magnitude) + assert layer.get_weight(torch.randn(out_dim, in_dim)).shape == (out_dim, in_dim) + + +def test_lokr_dora_magnitude_is_stripped(): + prefix = "diffusion_model.blocks.0.cross_attn.q_proj" + state_dict = { + f"{prefix}.lokr_w1": torch.zeros(2, 2), + f"{prefix}.lokr_w2": torch.zeros(3, 3), + f"{prefix}.magnitude": torch.ones(6), + } + + model = lora_model_from_anima_state_dict(state_dict) + layer = model.layers[f"{ANIMA_LORA_TRANSFORMER_PREFIX}blocks.0.cross_attn.q_proj"] + + assert isinstance(layer, LoKRLayer) diff --git a/tests/backend/patches/lora_conversions/test_flux_aitoolkit_lora_conversion_utils.py b/tests/backend/patches/lora_conversions/test_flux_aitoolkit_lora_conversion_utils.py index 648f17438bb..3941350c602 100644 --- a/tests/backend/patches/lora_conversions/test_flux_aitoolkit_lora_conversion_utils.py +++ b/tests/backend/patches/lora_conversions/test_flux_aitoolkit_lora_conversion_utils.py @@ -1,9 +1,11 @@ import accelerate import pytest +import torch from invokeai.backend.flux.model import Flux from invokeai.backend.flux.util import get_flux_transformers_params from invokeai.backend.model_manager.taxonomy import FluxVariantType +from invokeai.backend.patches.layers.dora_layer import DoRALayer from invokeai.backend.patches.lora_conversions.flux_aitoolkit_lora_conversion_utils import ( _group_state_by_submodel, is_state_dict_likely_in_flux_aitoolkit_format, @@ -64,3 +66,18 @@ def test_lora_model_from_flux_aitoolkit_state_dict(): state_dict = keys_to_mock_state_dict(flux_aitoolkit_state_dict_keys) assert lora_model_from_flux_aitoolkit_state_dict(state_dict) + + +@pytest.mark.parametrize("magnitude_suffix", ["magnitude", "lora_magnitude_vector.weight"]) +def test_lora_model_from_flux_aitoolkit_dora_magnitude(magnitude_suffix: str): + state_dict = { + "diffusion_model.double_blocks.0.img_attn.proj.lora_A.weight": torch.zeros(2, 5), + "diffusion_model.double_blocks.0.img_attn.proj.lora_B.weight": torch.zeros(7, 2), + f"diffusion_model.double_blocks.0.img_attn.proj.{magnitude_suffix}": torch.arange(1, 8, dtype=torch.float32), + } + + model = lora_model_from_flux_aitoolkit_state_dict(state_dict) + layer = model.layers["lora_transformer-double_blocks.0.img_attn.proj"] + + assert isinstance(layer, DoRALayer) + assert layer.magnitude_is_out_dim is True diff --git a/tests/backend/patches/lora_conversions/test_flux_bfl_peft_lora_conversion_utils.py b/tests/backend/patches/lora_conversions/test_flux_bfl_peft_lora_conversion_utils.py new file mode 100644 index 00000000000..46b82e82d5d --- /dev/null +++ b/tests/backend/patches/lora_conversions/test_flux_bfl_peft_lora_conversion_utils.py @@ -0,0 +1,54 @@ +import pytest +import torch + +from invokeai.backend.patches.layers.dora_layer import DoRALayer +from invokeai.backend.patches.lora_conversions.flux_bfl_peft_lora_conversion_utils import ( + is_state_dict_likely_in_flux_bfl_peft_format, + lora_model_from_flux2_bfl_peft_state_dict, + lora_model_from_flux_bfl_peft_state_dict, +) +from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_TRANSFORMER_PREFIX + + +@pytest.mark.parametrize("magnitude_suffix", ["magnitude", "lora_magnitude_vector.weight"]) +def test_dora_magnitude_is_preserved(magnitude_suffix: str): + in_dim, out_dim, rank = 5, 7, 2 + magnitude = torch.arange(1, out_dim + 1, dtype=torch.float32) + prefix = "diffusion_model.double_blocks.0.img_attn.proj" + state_dict = { + f"{prefix}.lora_A.weight": torch.zeros(rank, in_dim), + f"{prefix}.lora_B.weight": torch.zeros(out_dim, rank), + f"{prefix}.{magnitude_suffix}": magnitude, + } + + assert is_state_dict_likely_in_flux_bfl_peft_format(state_dict) + model = lora_model_from_flux_bfl_peft_state_dict(state_dict) + layer = model.layers[f"{FLUX_LORA_TRANSFORMER_PREFIX}double_blocks.0.img_attn.proj"] + + assert isinstance(layer, DoRALayer) + assert layer.magnitude_is_out_dim is True + assert torch.equal(layer.dora_scale, magnitude) + assert layer.get_weight(torch.randn(out_dim, in_dim)).shape == (out_dim, in_dim) + + +def test_flux2_fused_qkv_dora_magnitude_is_split_by_output(): + hidden_dim, rank = 6, 2 + magnitude = torch.arange(1, 3 * hidden_dim + 1, dtype=torch.float32) + prefix = "diffusion_model.double_blocks.0.img_attn.qkv" + state_dict = { + f"{prefix}.lora_A.weight": torch.zeros(rank, hidden_dim), + f"{prefix}.lora_B.weight": torch.zeros(3 * hidden_dim, rank), + f"{prefix}.magnitude": magnitude, + } + + model = lora_model_from_flux2_bfl_peft_state_dict(state_dict) + layer_names = ("to_q", "to_k", "to_v") + for index, layer_name in enumerate(layer_names): + key = f"{FLUX_LORA_TRANSFORMER_PREFIX}transformer_blocks.0.attn.{layer_name}" + layer = model.layers[key] + expected_magnitude = magnitude[index * hidden_dim : (index + 1) * hidden_dim] + + assert isinstance(layer, DoRALayer) + assert layer.magnitude_is_out_dim is True + assert torch.equal(layer.dora_scale, expected_magnitude) + assert layer.get_weight(torch.randn(hidden_dim, hidden_dim)).shape == (hidden_dim, hidden_dim) diff --git a/tests/backend/patches/lora_conversions/test_flux_diffusers_lora_conversion_utils.py b/tests/backend/patches/lora_conversions/test_flux_diffusers_lora_conversion_utils.py index be94504abc9..934d4d132ba 100644 --- a/tests/backend/patches/lora_conversions/test_flux_diffusers_lora_conversion_utils.py +++ b/tests/backend/patches/lora_conversions/test_flux_diffusers_lora_conversion_utils.py @@ -1,6 +1,7 @@ import pytest import torch +from invokeai.backend.patches.layers.dora_layer import DoRALayer from invokeai.backend.patches.lora_conversions.flux_diffusers_lora_conversion_utils import ( is_state_dict_likely_in_flux_diffusers_format, lora_model_from_flux_diffusers_state_dict, @@ -92,3 +93,19 @@ def test_lora_model_from_flux_diffusers_state_dict_extra_keys_error(): # Check that an error is raised. with pytest.raises(AssertionError): lora_model_from_flux_diffusers_state_dict(state_dict, alpha=8.0) + + +@pytest.mark.parametrize("magnitude_suffix", ["magnitude", "lora_magnitude_vector.weight"]) +def test_lora_model_from_flux_diffusers_dora_magnitude(magnitude_suffix: str): + prefix = "transformer.transformer_blocks.0.attn.to_out.0" + state_dict = { + f"{prefix}.lora_A.weight": torch.zeros(2, 5), + f"{prefix}.lora_B.weight": torch.zeros(7, 2), + f"{prefix}.{magnitude_suffix}": torch.arange(1, 8, dtype=torch.float32), + } + + model = lora_model_from_flux_diffusers_state_dict(state_dict, alpha=None) + layer = model.layers["lora_transformer-double_blocks.0.img_attn.proj"] + + assert isinstance(layer, DoRALayer) + assert layer.magnitude_is_out_dim is True diff --git a/tests/backend/patches/lora_conversions/test_flux_kohya_lora_conversion_utils.py b/tests/backend/patches/lora_conversions/test_flux_kohya_lora_conversion_utils.py index 35a5f5a9092..355d36b5767 100644 --- a/tests/backend/patches/lora_conversions/test_flux_kohya_lora_conversion_utils.py +++ b/tests/backend/patches/lora_conversions/test_flux_kohya_lora_conversion_utils.py @@ -5,6 +5,7 @@ from invokeai.backend.flux.model import Flux from invokeai.backend.flux.util import get_flux_transformers_params from invokeai.backend.model_manager.taxonomy import FluxVariantType +from invokeai.backend.patches.layers.dora_layer import DoRALayer from invokeai.backend.patches.lora_conversions.flux_kohya_lora_conversion_utils import ( _convert_flux_transformer_kohya_state_dict_to_invoke_format, is_state_dict_likely_in_flux_kohya_format, @@ -114,3 +115,17 @@ def test_lora_model_from_flux_kohya_state_dict(sd_keys: dict[str, list[int]]): lora_model_keys = set(lora_model.layers.keys()) lora_model_keys = {k.replace(".", "_") for k in lora_model_keys} assert lora_model_keys == expected_layer_keys + + +def test_lora_model_from_flux_kohya_dora_magnitude(): + state_dict = { + "lora_unet_double_blocks_0_img_attn_proj.lora_down.weight": torch.zeros(2, 5), + "lora_unet_double_blocks_0_img_attn_proj.lora_up.weight": torch.zeros(7, 2), + "lora_unet_double_blocks_0_img_attn_proj.magnitude": torch.arange(1, 8, dtype=torch.float32), + } + + model = lora_model_from_flux_kohya_state_dict(state_dict) + layer = model.layers["lora_transformer-double_blocks.0.img_attn.proj"] + + assert isinstance(layer, DoRALayer) + assert layer.magnitude_is_out_dim is True diff --git a/tests/backend/patches/lora_conversions/test_qwen_image_lora_conversion_utils.py b/tests/backend/patches/lora_conversions/test_qwen_image_lora_conversion_utils.py index 642a20b7164..1be31f95c10 100644 --- a/tests/backend/patches/lora_conversions/test_qwen_image_lora_conversion_utils.py +++ b/tests/backend/patches/lora_conversions/test_qwen_image_lora_conversion_utils.py @@ -1,7 +1,9 @@ """Tests for Qwen Image LoRA conversion utilities.""" +import pytest import torch +from invokeai.backend.patches.layers.dora_layer import DoRALayer from invokeai.backend.patches.lora_conversions.qwen_image_lora_constants import ( QWEN_IMAGE_EDIT_LORA_TRANSFORMER_PREFIX, ) @@ -117,6 +119,26 @@ def test_diffusers_with_transformer_prefix_strips_it(): assert expected_key in model.layers +@pytest.mark.parametrize("magnitude_suffix", ["magnitude", "lora_magnitude_vector.weight"]) +def test_diffusers_dora_magnitude_is_preserved(magnitude_suffix: str): + in_dim, out_dim, rank = 5, 7, 2 + magnitude = torch.arange(1, out_dim + 1, dtype=torch.float32) + prefix = "transformer.transformer_blocks.0.attn.to_q" + state_dict = { + f"{prefix}.lora_A.weight": torch.zeros(rank, in_dim), + f"{prefix}.lora_B.weight": torch.zeros(out_dim, rank), + f"{prefix}.{magnitude_suffix}": magnitude, + } + + model = lora_model_from_qwen_image_state_dict(state_dict) + layer = model.layers[f"{QWEN_IMAGE_EDIT_LORA_TRANSFORMER_PREFIX}transformer_blocks.0.attn.to_q"] + + assert isinstance(layer, DoRALayer) + assert layer.magnitude_is_out_dim is True + assert torch.equal(layer.dora_scale, magnitude) + assert layer.get_weight(torch.randn(out_dim, in_dim)).shape == (out_dim, in_dim) + + # ---- Unknown key handling tests ---- diff --git a/tests/backend/patches/lora_conversions/test_wan_lora_conversion_utils.py b/tests/backend/patches/lora_conversions/test_wan_lora_conversion_utils.py index f9cac4bd61b..827c3a4eab0 100644 --- a/tests/backend/patches/lora_conversions/test_wan_lora_conversion_utils.py +++ b/tests/backend/patches/lora_conversions/test_wan_lora_conversion_utils.py @@ -1,7 +1,9 @@ """Tests for Wan LoRA state-dict conversion to ModelPatchRaw.""" +import pytest import torch +from invokeai.backend.patches.layers.dora_layer import DoRALayer from invokeai.backend.patches.lora_conversions.wan_lora_constants import WAN_LORA_TRANSFORMER_PREFIX from invokeai.backend.patches.lora_conversions.wan_lora_conversion_utils import ( _kohya_layer_to_diffusers_path, @@ -173,3 +175,22 @@ def test_unknown_kohya_submodule_is_skipped_silently(self): def test_empty_state_dict(self): patch = lora_model_from_wan_state_dict({}) assert len(patch.layers) == 0 + + @pytest.mark.parametrize("magnitude_suffix", ["magnitude", "lora_magnitude_vector.weight"]) + def test_peft_dora_magnitude_is_preserved(self, magnitude_suffix: str): + in_dim, out_dim, rank = 5, 7, 2 + magnitude = torch.arange(1, out_dim + 1, dtype=torch.float32) + prefix = "diffusion_model.blocks.0.self_attn.q" + state_dict = { + f"{prefix}.lora_A.weight": torch.zeros(rank, in_dim), + f"{prefix}.lora_B.weight": torch.zeros(out_dim, rank), + f"{prefix}.{magnitude_suffix}": magnitude, + } + + patch = lora_model_from_wan_state_dict(state_dict) + layer = patch.layers[f"{WAN_LORA_TRANSFORMER_PREFIX}blocks.0.attn1.to_q"] + + assert isinstance(layer, DoRALayer) + assert layer.magnitude_is_out_dim is True + assert torch.equal(layer.dora_scale, magnitude) + assert layer.get_weight(torch.randn(out_dim, in_dim)).shape == (out_dim, in_dim) diff --git a/tests/backend/patches/lora_conversions/test_z_image_lora_conversion_utils.py b/tests/backend/patches/lora_conversions/test_z_image_lora_conversion_utils.py new file mode 100644 index 00000000000..e03b6b57536 --- /dev/null +++ b/tests/backend/patches/lora_conversions/test_z_image_lora_conversion_utils.py @@ -0,0 +1,28 @@ +import pytest +import torch + +from invokeai.backend.patches.layers.dora_layer import DoRALayer +from invokeai.backend.patches.lora_conversions.z_image_lora_constants import Z_IMAGE_LORA_TRANSFORMER_PREFIX +from invokeai.backend.patches.lora_conversions.z_image_lora_conversion_utils import ( + lora_model_from_z_image_state_dict, +) + + +@pytest.mark.parametrize("magnitude_suffix", ["magnitude", "lora_magnitude_vector.weight"]) +def test_peft_dora_magnitude_is_preserved(magnitude_suffix: str): + in_dim, out_dim, rank = 5, 7, 2 + magnitude = torch.arange(1, out_dim + 1, dtype=torch.float32) + prefix = "diffusion_model.layers.0.attention.to_q" + state_dict = { + f"{prefix}.lora_A.weight": torch.zeros(rank, in_dim), + f"{prefix}.lora_B.weight": torch.zeros(out_dim, rank), + f"{prefix}.{magnitude_suffix}": magnitude, + } + + model = lora_model_from_z_image_state_dict(state_dict) + layer = model.layers[f"{Z_IMAGE_LORA_TRANSFORMER_PREFIX}layers.0.attention.to_q"] + + assert isinstance(layer, DoRALayer) + assert layer.magnitude_is_out_dim is True + assert torch.equal(layer.dora_scale, magnitude) + assert layer.get_weight(torch.randn(out_dim, in_dim)).shape == (out_dim, in_dim)