invoke-ai / invoke-ai/InvokeAI

ai-toolkit DoRA (.magnitude) adapters fail to load for every base except Krea-2

Open
#9,558 0 comments 0 reactions 1 assignee Claimed by @Pfannkuchensack View on GitHub
Dominant language
Python
Stars
28.2k
Forks
3k
Avg merge
6d 5h
Merged PRs (30d)
19

Description

### Summary

#9517 taught the **Krea-2** LoRA converter that ai-toolkit writes the DoRA magnitude as a bare `.magnitude` suffix (PEFT/out-dim orientation), and gave `DoRALayer` an explicit `magnitude_is_out_dim` flag. Every other architecture's converter still has a fixed suffix table that doesn't know `.magnitude`, so ai-toolkit DoRA adapters for those bases fail to load with the same error #9517 fixed for Krea-2:

```
ValueError: Unsupported lora format: dict_keys(['to_gate.magnitude', 'to_k.magnitude', ...])
```

### Why it happens

`.magnitude` isn't in the converter's known-suffix list, so `_group_by_layer` falls through to a blind `rsplit(".", 2)`. That cuts inside the module path and fuses every magnitude in a block into one bogus parent layer, which then matches no layer type.

Affected suffix tables (all on `main`):

- `invokeai/backend/patches/lora_conversions/wan_lora_conversion_utils.py:223-228`
- `invokeai/backend/patches/lora_conversions/z_image_lora_conversion_utils.py:229-234`
- `invokeai/backend/patches/lora_conversions/anima_lora_conversion_utils.py:152-157`
- `invokeai/backend/patches/lora_conversions/qwen_image_lora_conversion_utils.py:168-173`
- `invokeai/backend/patches/lora_conversions/flux_bfl_peft_lora_conversion_utils.py:78` uses a similar fixed suffix set and is worth checking too

None of these map `.lora_magnitude_vector.weight` either, so unlike Krea-2 they never produced a *mis-oriented* DoRA layer — this is a plain "fails to load", not a silent-wrong-weights bug. That's why it's a follow-up rather than part of #9517.

### Suggested fix

The machinery already exists after #9517 — this is mostly plumbing:

1. Add `.magnitude` (and `.lora_magnitude_vector.weight`) to each converter's suffix table, routing **both to the `dora_magnitude` value key**, not `dora_scale`. The two conventions index opposite axes of the `(out_features, in_features)` weight and must not be mixed up — see the `DoRALayer` docstring added in #9517.
2. `any_lora_layer_from_state_dict` already dispatches on `dora_magnitude`, so no change is needed there.
3. Watch the LoKr interaction: `anima_lora_conversion_utils._make_layer_patch` strips only `dora_scale` from LoKr layers before dispatch. If Anima learns `.magnitude`, that strip needs to cover `dora_magnitude` as well, or a DoRA+LoKr adapter will take the DoRA branch and die on a missing `lora_up.weight`.
4. Per-architecture verification is worth doing rather than assuming: instantiate the target transformer on the `meta` device and confirm every converted key resolves to a real module with `magnitude.numel() == out_features`. That's the check that gave #9517 its 256/256 clean result.

### Notes

Non-blocking follow-up raised during the review of #9517 (https://github.com/invoke-ai/InvokeAI/pull/9517#pullrequestreview-5058290795), finding 4. No user has reported this for a specific non-Krea-2 base yet, so priority should probably follow demand — ai-toolkit is a popular trainer, so it's likely only a matter of time.

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.