Skip to content

fix(lora): handle DoRA magnitudes on LoKr and LoHa layers - #10075

Merged
lstein merged 5 commits into
invoke-ai:mainfrom
guishiw:fix/10018-dokr-doha-fallback
Nov 7, 2026
Merged

lstein merged 5 commits into
invoke-ai:mainfrom
guishiw:fix/10018-dokr-doha-fallback

Conversation

@guishiw

@guishiw guishiw commented Oct 30, 2026 •

Copy link
Copy Markdown
Contributor

Summary

Fixes #10018

Root cause

The shared layer factory dispatched any state dict containing dora_scale or dora_magnitude to DoRALayer before checking for LoKr or LoHa factors. DoKr and DoHa layers therefore tried to read nonexistent lora_up.weight / lora_down.weight tensors and failed with KeyError.

Fix strategy

Detect LoKr and LoHa factorized layouts before plain DoRA. Because combined DoRA magnitude application is not supported for these factorized layers, drop only the magnitude tensor, apply the plain LoKr/LoHa delta, and emit one warning. The shared handling replaces the Anima- and Krea-2-specific LoKr workarounds.

Related Issues / Discussions

Closes #10018

QA Instructions

  • pytest tests/backend/patches/layers/test_utils.py tests/backend/patches/lora_conversions/test_anima_lora_conversion_utils.py tests/backend/patches/lora_conversions/test_krea2_lora_conversion_utils.py -q
    • 121 passed
  • Ruff check and format --check pass on all changed Python files.

The relevant test suite passes with 121 tests, covering both magnitude key spellings for LoKr and LoHa plus the existing Anima and Krea-2 conversion paths.

Review

Plain LoRA+DoRA dispatch is unchanged. The fallback is limited to state dicts that already contain LoKr or LoHa structural factors, and regression assertions verify the resulting factorized weight and single warning.

Compatibility / Rollout

No persisted-state, API, or rollout changes. Unsupported DoKr/DoHa magnitudes now degrade explicitly to plain LoKr/LoHa instead of crashing.

Checklist

  • The PR has a short but descriptive title, suitable for a changelog
  • Meaningful regression coverage added / updated where needed; obsolete tests/code removed
  • Persisted-state and API changes include required migrations / compatibility validation
  • Relevant performance/efficiency opportunities considered; material claims have evidence
  • Material review findings resolved and relevant checks rerun
  • Documentation added / updated (if applicable)
  • Updated What's New copy (if doing a release after this PR)

View with [code]smith Autofix with [code]smith
Need help on this PR? Tag @codesmith-bot with what you need.

@github-actions github-actions Bot added python PRs that change python files backend PRs that change backend files python-tests PRs that change python tests labels Oct 30, 2026

@Pfannkuchensack Pfannkuchensack left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, this fixes #10018: LoKr and LoHa layers with a DoRA magnitude load on every base now and apply as plain LoKr/LoHa, and the Anima and Krea-2 workarounds are gone. Checked:

  • Real file from #10018. Car Meet Garage (SDXL, LyCORIS algo=lokr, dora_wd=True, dora_scale on every module) goes through probe (LoRA_LyCORIS_SDXL_Config, SDXL, LyCORIS), LoRALoader and module resolution on meta-device SDXL models:
    • main fails with KeyError: 'lora_up.weight'.
    • Here: 1052 LoKRLayers (912 with lokr_w1 + lokr_w2_a/b, 140 with lokr_w1 + lokr_w2, 49 of them Conv2d). All resolve: 788 on the UNet, 72 on CLIP-L, 192 on CLIP-G. No unresolved layer, no shape mismatch.
    • Against the real SDXL base 1.0 weights, every layer's delta is exactly alpha / rank * kron(w1, w2) (1052/1052). Alpha is applied once: 0.5 on the factored layers, 1.0 on the unfactored ones, the same as LyCORIS (it forces alpha = lora_dim when both operands are whole).
    • On an RTX 4090 (SDXL base 1.0, 30 steps, 1024²) the LoRA generates with a clearly visible effect (car style, colour and scene change against no LoRA; weight 0.6 lands in between). Two runs with the same seed are bit-identical.
  • Every variant. Synthetic LoKr (whole and factored operands, Linear, Conv2d with a 4-D w2, with a flattened w2_a/b, and Tucker t2) and LoHa (Linear, flattened conv, Tucker t1/t2), each with dora_scale on the output axis ([out, 1]), on the input axis ([1, in]), and with dora_magnitude. That makes 27 cases. All fail on main with the KeyError. Here all of them load as the right layer type, and the delta equals alpha / rank times an independent rebuild of the factors.
  • Anima and Krea-2. Both converters already dropped the magnitude from LoKr layers and applied plain LoKr. The shared factory now does the same, so their LoKr behaviour is unchanged: same layer, same weights, same alpha. There is no DoRA axis involved because the magnitude is discarded. The only differences: a LoHa layer with a magnitude, and on Anima a LoKr layer with dora_magnitude, used to crash and now load as plain. The Krea-2 warning also loses the layer key. LoRA+DoRA layers still reach DoRALayer with the axis detection from #10039: the Krea-2 DoRA tests and test_dora_layer.py pass.
  • Other layer types. Moving the LoHa and LoKr checks ahead of the DoRA check only changes dispatch for state dicts that contain hada_w1_a, lokr_w1 or lokr_w1_a. Plain LoRA, LoRA+DoRA, full, IA3, norm and sparse bias dispatch as before.
  • Tests. The new tests fail on main for the right reason (4 × KeyError: 'lora_up.weight'). Your QA command: 121 passed. tests/backend/patches: 310 passed. tests/backend/patches + tests/backend/model_manager (-n 4): 2638 passed and 9 failed, all of them environmental. 8 are test_huggingface_metadata_fetch_timeout.py, which fails the same way on main (a local HF_ENDPOINT proxy). The 9th is test_16_channel_vae_loader, which reads an LFS pointer in this checkout. Ruff check and format are clean.

Blocking:

  • One warning per layer instead of one per LoRA (invokeai/backend/patches/layers/utils.py:18-27). _without_unsupported_magnitude logs every time the factory sees a magnitude. Loading the #10018 file once logs "LoKr layer combines factorized weights with a DoRA magnitude, which is not supported; applying it as plain LoKr." 1052 times (788 UNet + 264 text encoder layers). I counted that on CPU through LoRALoader._load_model, and the same count shows up on a server running this branch. The warning doesn't name the LoRA, so with two such LoRAs in a stack the log can't say which one it was. The PR description and #10018 both ask for a single warning. Anima (_make_layer_patch) and Krea-2 (_get_lora_layer_values) already warned per layer, but only for those two bases. This PR spreads that behaviour to every base, and SDXL/Illustrious DoKr files are common on Civitai. Please warn once per LoRA and name it. For example, the factory could drop the magnitude silently and LoRALoader._load_model could warn once with config.name and the number of affected layers. Or the factory could report the drop back to its callers. The new tests pin the per-layer message (tests/backend/patches/layers/test_lora_layer_factory.py:24, :48), so they need to move along with it: a two-layer file should give exactly one warning.

Some points, none blocking:

  • How much the fallback drops. #10018 allows plain LoKr/LoHa, so this is fine for this PR. The dropped part isn't small, though, so DoKr/DoHa would be a worthwhile follow-up. On the real file with the real SDXL base weights, dora_scale differs from the row norm of W + ΔW_LoKr by a median 0.33 % per UNet layer (p90 0.35 %, max 0.92 %), and by up to 4.4 % on single channels (UNet row factor 0.956 to 1.024, worst on up_blocks.1 … attn2.to_out.0). About 0.1 % of that comes from bf16 rounding of dora_scale. The LoKr delta itself is only a median 0.26 % of ‖W‖, while the full DoRA update is 0.45 %. As a result, the update this PR applies differs from the LyCORIS DoRA update by a median 82 % of that update's norm on the UNet (p90 89 %), 43 % on CLIP-L and 72 % on CLIP-G. ComfyUI (comfy/weight_adapter/base.py:weight_decompose, used by its LoKr and LoHa adapters) and LyCORIS apply the magnitude. The DoRA math doesn't depend on how ΔW is factored, so a follow-up could reuse DoRALayer's axis detection (_dora_scale_is_out_dim) around any LoKr/LoHa delta. I'll open an issue for that unless you want to take it on.
  • Pass-through left behind (anima_lora_conversion_utils.py:148). _make_layer_patch now only calls any_lora_layer_from_state_dict. Calling the factory directly at its two call sites (:261, :299) removes the wrapper.
  • Stale comments in the converter tests. test_anima_lora_conversion_utils.py:287-291 ("should be stripped … since shared code doesn't support") and test_krea2_lora_conversion_utils.py:387-388 and :561-562 ("The converter drops it", "Kept, it would dispatch the layer to DoRALayer") describe the removed workarounds. The shared factory does this now.
  • Test strength. Neither new test checks scaling. The LoKr case has two whole operands, so its scale is 1, and the LoHa case has alpha == rank == 2. Both compare get_weight, which leaves the scale out. With scale() patched to return its square, or a constant 1.0, the new tests and the Anima DoRA test still pass. Using a factored LoKr operand and alpha != rank, and comparing get_parameters(..., 1.0)["weight"], would cover alpha.
  • Docs. model-format-support.mdx lists LoKr and LoHa support per base but says nothing about DoRA magnitudes. One sentence saying that LoKr/LoHa files with a DoRA magnitude load as plain LoKr/LoHa would tell users why such a LoRA can look weaker than in ComfyUI.

@guishiw

guishiw commented Oct 30, 2026

Copy link
Copy Markdown
Contributor Author
  • One warning per layer instead of one per LoRA

Thanks for the review. I agree that logging once per layer is too noisy for real-world LoRAs. I’ll update this to emit a single warning per LoRA, including the LoRA name and the number of affected layers, and add a regression test covering multiple affected layers while asserting that only one warning is emitted.

I’ll keep the broader DoKr/DoHa support and the non-blocking cleanup/documentation suggestions as follow-up work.

@Pfannkuchensack Pfannkuchensack left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, c818e18 resolves the blocking point. Rechecked on 189c516:

  • Real file from #10018. Car Meet Garage through LoRALoader._load_model now logs exactly one warning: LoRA 'sdxl_lycoris-lokr-dora_car-meet-garage-1967-stingray' contains 1052 LoKr/LoHa layers with unsupported DoRA magnitudes; applying those layers as plain LoKr/LoHa. I counted every warning that reached any logger, not only the loader's. The file still gives 1052 LoKRLayers, with the same alpha scales as before (0.5 on 912 layers, 1.0 on 140).
  • Other layouts. Two-layer LoKr and LoHa files on SD 1.5, SDXL, FLUX (kohya and BFL keys), Qwen-Image (kohya keys), Anima and Krea-2: every layout that loads without a magnitude also loads with dora_scale and warns once. The warning counts layers correctly for each.
  • Checks. tests/backend/patches + tests/backend/model_manager/load + tests/backend/model_manager/configs: 2805 passed, 152 skipped, 1 failed. The failure is test_16_channel_vae_loader, which reads an LFS pointer in my checkout. Ruff is clean, and CI is green.

Two small points, neither blocking:

  • The new test can't catch the original regression. test_factorized_dora_layers_warn_once_per_lora_with_name_and_count asserts on the loader's mock logger, while the per-layer warning came from the module logger in layers/utils.py. With that warning restored, this test would still pass, and the caplog assertions that used to see it are gone. assert caplog.messages == [] in the two factory tests would cover it.
  • Two of the counted suffixes aren't dropped by the factory. _DORA_MAGNITUDE_SUFFIXES includes .magnitude and .lora_magnitude_vector.weight, but _without_unsupported_magnitude only removes dora_scale and dora_magnitude. Krea-2 renames both first, so it stays at one warning. On kohya keys they reach LoKRLayer/LoHALayer unrenamed, and each layer then logs Unexpected keys found. I only got there with synthetic files. LyCORIS writes dora_scale, so this shouldn't show up in practice.

Dotted LoKr on Z-Image and Wan, and dotted LoHa on Z-Image, Wan, Anima and Krea-2, still fail with or without a magnitude. That's the same gap as #10020, not something this PR changes. Applying the magnitude is tracked in #10107.

Approving.

@lstein
lstein enabled auto-merge (squash) November 7, 2026 16:17
@lstein
lstein merged commit a86f0dd into invoke-ai:main Nov 7, 2026
20 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

7.0.0 backend PRs that change backend files python PRs that change python files python-tests PRs that change python tests

Projects

Status: 7.0.0: UX overhaul

Development

Successfully merging this pull request may close these issues.

[bug]: LoKr and LoHa layers with a DoRA magnitude fail to load (KeyError: 'lora_up.weight')

3 participants