diff --git a/invokeai/backend/patches/layers/dora_layer.py b/invokeai/backend/patches/layers/dora_layer.py index 3e52ce95783..e2b16a6e2c0 100644 --- a/invokeai/backend/patches/layers/dora_layer.py +++ b/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 diff --git a/invokeai/backend/patches/layers/utils.py b/invokeai/backend/patches/layers/utils.py index 8141a56644a..e6d36e2f7c9 100644 --- a/invokeai/backend/patches/layers/utils.py +++ b/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 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 458640c906d..2605c18bdcb 100644 --- a/invokeai/backend/patches/lora_conversions/krea2_lora_conversion_utils.py +++ b/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 ``.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", } diff --git a/tests/backend/patches/layers/test_dora_layer.py b/tests/backend/patches/layers/test_dora_layer.py new file mode 100644 index 00000000000..a56afacf591 --- /dev/null +++ b/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) 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 d4ebcc7a431..f2953e47f40 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 @@ -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: