fix(lora): load ai-toolkit DoRA magnitudes across architectures

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

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",

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

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]

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.

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.

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.

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,13 +147,21 @@ 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"],
             "lora_up.weight": layer_dict["lora_B.weight"],
         }
         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,25 +173,27 @@ 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):
             continue
 
         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:

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,25 +221,27 @@ 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):
             continue
 
         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: