fix(lora): support ai-toolkit DoRA magnitudes for Krea-2 LoRAs

invokeai/backend/patches/layers/dora_layer.py

@@ -8,7 +8,21 @@
 
 
 class DoRALayer(LoRALayerBase):
-    """A DoRA layer. As defined in https://arxiv.org/pdf/2402.09353."""
+    """A DoRA layer. As defined in https://arxiv.org/pdf/2402.09353.
+
+    Two mutually-incompatible magnitude conventions exist in the wild, and they differ in *which axis* of the
+    ``(out_features, in_features)`` weight the magnitude vector indexes:
+
+    - **LyCORIS / kohya / OneTrainer** (``.dora_scale``): one entry per **input** column, i.e. the direction
+      norm is taken across the output dim. Shape ``(1, in_features)``.
+    - **PEFT / Diffusers / ai-toolkit** (``.lora_magnitude_vector.weight`` / ``.magnitude``): one entry per
+      **output** row, i.e. the direction norm is taken across the input dim (``torch.linalg.norm(w, dim=1)``).
+      Shape ``(out_features,)``.
+
+    Applying one convention's magnitude with the other's math silently produces wrong weights on square layers
+    and raises a broadcast error on non-square ones, so the orientation is carried explicitly rather than
+    guessed from the tensor shape.
+    """
 
     def __init__(
         self,
@@ -17,11 +31,15 @@ def __init__(
         dora_scale: torch.Tensor,
         alpha: float | None,
         bias: Optional[torch.Tensor],
+        magnitude_is_out_dim: bool = False,
     ):
         super().__init__(alpha, bias)
         self.up = up
         self.down = down
         self.dora_scale = dora_scale
+        # False -> LyCORIS convention (magnitude indexes in_features). True -> PEFT/ai-toolkit convention
+        # (magnitude indexes out_features).
+        self.magnitude_is_out_dim = magnitude_is_out_dim
 
     @classmethod
     def from_state_dict_values(cls, values: Dict[str, torch.Tensor]):
@@ -30,12 +48,17 @@ def from_state_dict_values(cls, values: Dict[str, torch.Tensor]):
             values.get("bias_indices", None), values.get("bias_values", None), values.get("bias_size", None)
         )
 
+        # ``dora_magnitude`` is the PEFT/ai-toolkit out-dim magnitude; ``dora_scale`` is the LyCORIS in-dim one.
+        magnitude_is_out_dim = "dora_magnitude" in values
+        dora_scale = values["dora_magnitude"] if magnitude_is_out_dim else values["dora_scale"]
+
         layer = cls(
             up=values["lora_up.weight"],
             down=values["lora_down.weight"],
-            dora_scale=values["dora_scale"],
+            dora_scale=dora_scale,
             alpha=alpha,
             bias=bias,
+            magnitude_is_out_dim=magnitude_is_out_dim,
         )
 
         cls.warn_on_unhandled_keys(
@@ -50,6 +73,7 @@ def from_state_dict_values(cls, values: Dict[str, torch.Tensor]):
                 "lora_up.weight",
                 "lora_down.weight",
                 "dora_scale",
+                "dora_magnitude",
             },
         )
 
@@ -70,16 +94,26 @@ def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
         # At this point, out_weight is the unnormalized direction matrix.
         out_weight = orig_weight + delta_v
 
-        # TODO(ryand): Simplify this logic.
-        direction_norm = (
-            out_weight.transpose(0, 1)
-            .reshape(out_weight.shape[1], -1)
-            .norm(dim=1, keepdim=True)
-            .reshape(out_weight.shape[1], *[1] * (out_weight.dim() - 1))
-            .transpose(0, 1)
-        )
+        if self.magnitude_is_out_dim:
+            # PEFT / ai-toolkit: norm over the input dim, one entry per output row.
+            trailing_dims = [1] * (out_weight.dim() - 1)
+            direction_norm = (
+                out_weight.reshape(out_weight.shape[0], -1).norm(dim=1).reshape(out_weight.shape[0], *trailing_dims)
+            )
+            dora_scale = self.dora_scale.reshape(out_weight.shape[0], *trailing_dims)
+        else:
+            # LyCORIS / kohya: norm over the output dim, one entry per input column.
+            # TODO(ryand): Simplify this logic.
+            direction_norm = (
+                out_weight.transpose(0, 1)
+                .reshape(out_weight.shape[1], -1)
+                .norm(dim=1, keepdim=True)
+                .reshape(out_weight.shape[1], *[1] * (out_weight.dim() - 1))
+                .transpose(0, 1)
+            )
+            dora_scale = self.dora_scale
 
-        out_weight *= self.dora_scale / direction_norm
+        out_weight *= dora_scale / direction_norm
 
         return out_weight - orig_weight
 

invokeai/backend/patches/layers/utils.py

@@ -15,7 +15,8 @@
 def any_lora_layer_from_state_dict(state_dict: Dict[str, torch.Tensor]) -> BaseLayerPatch:
     # Detect layers according to LyCORIS detection logic(`weight_list_det`)
     # https://github.com/KohakuBlueleaf/LyCORIS/tree/8ad8000efb79e2b879054da8c9356e6143591bad/lycoris/modules
-    if "dora_scale" in state_dict:
+    if "dora_scale" in state_dict or "dora_magnitude" in state_dict:
+        # ``dora_scale`` is the LyCORIS in-dim magnitude, ``dora_magnitude`` the PEFT/ai-toolkit out-dim one.
         return DoRALayer.from_state_dict_values(state_dict)
     elif "lora_up.weight" in state_dict:
         # LoRA a.k.a LoCon

invokeai/backend/patches/lora_conversions/krea2_lora_conversion_utils.py

@@ -188,8 +188,9 @@ def _get_lora_layer_values(
             "lora_down.weight": layer_dict["lora_A.weight"],
             "lora_up.weight": layer_dict["lora_B.weight"],
         }
-        if "dora_scale" in layer_dict:
-            values["dora_scale"] = layer_dict["dora_scale"]
+        for magnitude_key in ("dora_scale", "dora_magnitude"):
+            if magnitude_key in layer_dict:
+                values[magnitude_key] = layer_dict[magnitude_key]
         if "alpha" in layer_dict:
             values["alpha"] = layer_dict["alpha"]
         if alpha is not None:
@@ -198,17 +199,22 @@ def _get_lora_layer_values(
     return layer_dict
 
 
-# Maps each recognized weight-key suffix to the canonical value-key used downstream. The PEFT/diffusers DoRA
-# magnitude is published as ``<layer>.lora_magnitude_vector.weight``; it is the same thing InvokeAI stores as
-# ``dora_scale``, so mapping it here lets a standard Diffusers DoRA adapter (A/B + magnitude) load as a
-# DoRALayer instead of being split into a bogus, unrecognized layer.
+# Maps each recognized weight-key suffix to the canonical value-key used downstream.
+#
+# DoRA magnitudes come in two orientations that must not be mixed up (see ``DoRALayer``):
+#   - ``.dora_scale`` (LyCORIS/kohya) indexes the *input* dim  -> value key ``dora_scale``
+#   - ``.lora_magnitude_vector.weight`` (PEFT/diffusers) and ``.magnitude`` (ai-toolkit) index the *output*
+#     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.
 _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_scale",
+    ".lora_magnitude_vector.weight": "dora_magnitude",
+    ".magnitude": "dora_magnitude",
     ".alpha": "alpha",
 }
 

tests/backend/patches/layers/test_dora_layer.py

@@ -0,0 +1,65 @@
+import torch
+
+from invokeai.backend.patches.layers.dora_layer import DoRALayer
+
+
+def _reference_aitoolkit_forward(
+    x: torch.Tensor, orig_weight: torch.Tensor, down: torch.Tensor, up: torch.Tensor, magnitude: torch.Tensor
+) -> torch.Tensor:
+    """The output ai-toolkit's DoRAModule produces (multiplier=1, alpha=rank so scale=1).
+
+    See ``ToolkitModuleMixin.forward`` + ``DoRAModule.apply_dora`` in ostris/ai-toolkit: the module output is
+    ``org_forward(x) + lora_output + (magnitude / ||W + dV||_row - 1) * F.linear(x, W + dV)``.
+    """
+    delta_v = up @ down
+    weight_norm = torch.linalg.norm(orig_weight + delta_v, dim=1)
+    return (
+        torch.nn.functional.linear(x, orig_weight)
+        + torch.nn.functional.linear(x, delta_v)
+        + (magnitude / weight_norm).view(1, -1) * torch.nn.functional.linear(x, orig_weight + delta_v)
+        - torch.nn.functional.linear(x, orig_weight + delta_v)
+    )
+
+
+@torch.no_grad()
+def test_out_dim_magnitude_matches_peft_aitoolkit_math() -> None:
+    """A PEFT/ai-toolkit DoRA magnitude (one entry per output row) must reproduce their forward pass.
+
+    Covers non-square layers in both directions: applying the LyCORIS (input-dim) math to an output-dim
+    magnitude raises a broadcast error there, and is silently wrong on square layers.
+    """
+    torch.manual_seed(0)
+    for out_features, in_features, rank in [(12, 20, 4), (16, 16, 4), (20, 12, 4)]:
+        orig_weight = torch.randn(out_features, in_features)
+        down = torch.randn(rank, in_features) * 0.05
+        up = torch.randn(out_features, rank) * 0.05
+        magnitude = torch.randn(out_features).abs() + 1.0
+
+        layer = DoRALayer.from_state_dict_values(
+            {"lora_down.weight": down, "lora_up.weight": up, "dora_magnitude": magnitude}
+        )
+        assert layer.magnitude_is_out_dim is True
+
+        patched_weight = orig_weight + layer.get_weight(orig_weight)
+        x = torch.randn(7, in_features)
+        expected = _reference_aitoolkit_forward(x, orig_weight, down, up, magnitude)
+        assert torch.allclose(torch.nn.functional.linear(x, patched_weight), expected, atol=1e-5)
+
+
+@torch.no_grad()
+def test_in_dim_magnitude_keeps_lycoris_math() -> None:
+    """The LyCORIS/kohya ``dora_scale`` (one entry per input column) keeps its original normalization."""
+    torch.manual_seed(0)
+    out_features, in_features, rank = 12, 20, 4
+    orig_weight = torch.randn(out_features, in_features)
+    down = torch.randn(rank, in_features) * 0.05
+    up = torch.randn(out_features, rank) * 0.05
+    dora_scale = torch.randn(1, in_features).abs() + 1.0
+
+    layer = DoRALayer.from_state_dict_values({"lora_down.weight": down, "lora_up.weight": up, "dora_scale": dora_scale})
+    assert layer.magnitude_is_out_dim is False
+
+    direction = orig_weight + up @ down
+    direction_norm = direction.transpose(0, 1).norm(dim=1, keepdim=True).transpose(0, 1)
+    expected = direction * (dora_scale / direction_norm)
+    assert torch.allclose(orig_weight + layer.get_weight(orig_weight), expected, atol=1e-6)

tests/backend/patches/lora_conversions/test_krea2_lora_conversion_utils.py

@@ -39,6 +39,8 @@ def test_peft_dora_layer_preserves_magnitude_and_alpha() -> None:
     assert isinstance(layer, DoRALayer)
     assert layer._alpha == 1.0
     assert torch.equal(layer.dora_scale, dora_scale)
+    # `.dora_scale` is the LyCORIS magnitude: it indexes the *input* dim.
+    assert layer.magnitude_is_out_dim is False
 
 
 def test_peft_layer_without_explicit_alpha_uses_rank_default() -> None:
@@ -85,6 +87,40 @@ def test_peft_dora_magnitude_vector_key_produces_dora_layer() -> None:
     layer = model.layers[f"{KREA2_LORA_TRANSFORMER_PREFIX}text_fusion.0.attn.to_q"]
     assert isinstance(layer, DoRALayer)
     assert torch.equal(layer.dora_scale, magnitude)
+    # The PEFT magnitude indexes the *output* dim, unlike the LyCORIS `.dora_scale`.
+    assert layer.magnitude_is_out_dim is True
+
+
+def test_native_aitoolkit_dora_magnitude_key_produces_dora_layer() -> None:
+    # ai-toolkit (`network.type: dora`) writes native Krea-2 keys with a bare `.magnitude` suffix. Without an
+    # explicit mapping these fall through the suffix table and get grouped into a bogus `...attn` layer,
+    # raising "Unsupported lora format: dict_keys(['to_gate.magnitude', ...])" (issue #9515).
+    attn_magnitude = torch.full((4,), 3.0)
+    # A non-square layer: its magnitude has out_features entries while the LyCORIS convention would expect
+    # in_features, so a mis-oriented magnitude would blow up at patch time rather than silently.
+    ff_magnitude = torch.full((4,), 5.0)
+    state_dict = {
+        "diffusion_model.blocks.0.attn.wq.lora_A.weight": torch.ones(2, 4),
+        "diffusion_model.blocks.0.attn.wq.lora_B.weight": torch.ones(4, 2),
+        "diffusion_model.blocks.0.attn.wq.magnitude": attn_magnitude,
+        "diffusion_model.txtfusion.refiner_blocks.0.mlp.down.lora_A.weight": torch.ones(2, 8),
+        "diffusion_model.txtfusion.refiner_blocks.0.mlp.down.lora_B.weight": torch.ones(4, 2),
+        "diffusion_model.txtfusion.refiner_blocks.0.mlp.down.magnitude": ff_magnitude,
+    }
+
+    model = lora_model_from_krea2_state_dict(state_dict)
+
+    attn_layer = model.layers[f"{KREA2_LORA_TRANSFORMER_PREFIX}transformer_blocks.0.attn.to_q"]
+    assert isinstance(attn_layer, DoRALayer)
+    assert torch.equal(attn_layer.dora_scale, attn_magnitude)
+    assert attn_layer.magnitude_is_out_dim is True
+
+    ff_layer = model.layers[f"{KREA2_LORA_TRANSFORMER_PREFIX}text_fusion.refiner_blocks.0.ff.down"]
+    assert isinstance(ff_layer, DoRALayer)
+    assert torch.equal(ff_layer.dora_scale, ff_magnitude)
+    assert ff_layer.magnitude_is_out_dim is True
+    # The magnitude must survive as a real DoRA layer, not leak into a bogus parent group.
+    assert not any(key.endswith(".attn") or key.endswith(".mlp") for key in model.layers)
 
 
 def test_conflicting_transformer_and_diffusion_model_aliases_raise() -> None: