Repository navigation
fix(lora): handle DoRA magnitudes on LoKr and LoHa layers - #10075
Conversation
Pfannkuchensack
left a comment
There was a problem hiding this comment.
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_scaleon every module) goes through probe (LoRA_LyCORIS_SDXL_Config, SDXL, LyCORIS),LoRALoaderand module resolution on meta-device SDXL models:mainfails withKeyError: 'lora_up.weight'.- Here: 1052
LoKRLayers (912 withlokr_w1+lokr_w2_a/b, 140 withlokr_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 forcesalpha = lora_dimwhen 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 flattenedw2_a/b, and Tuckert2) and LoHa (Linear, flattened conv, Tuckert1/t2), each withdora_scaleon the output axis ([out, 1]), on the input axis ([1, in]), and withdora_magnitude. That makes 27 cases. All fail onmainwith theKeyError. Here all of them load as the right layer type, and the delta equalsalpha / ranktimes 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 reachDoRALayerwith the axis detection from #10039: the Krea-2 DoRA tests andtest_dora_layer.pypass. - 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_w1orlokr_w1_a. Plain LoRA, LoRA+DoRA, full, IA3, norm and sparse bias dispatch as before. - Tests. The new tests fail on
mainfor 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 aretest_huggingface_metadata_fetch_timeout.py, which fails the same way onmain(a localHF_ENDPOINTproxy). The 9th istest_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_magnitudelogs 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 throughLoRALoader._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 andLoRALoader._load_modelcould warn once withconfig.nameand 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_scalediffers from the row norm ofW + ΔW_LoKrby 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 onup_blocks.1 … attn2.to_out.0). About 0.1 % of that comes from bf16 rounding ofdora_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 reuseDoRALayer'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_patchnow only callsany_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") andtest_krea2_lora_conversion_utils.py:387-388and: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 compareget_weight, which leaves the scale out. Withscale()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 andalpha != rank, and comparingget_parameters(..., 1.0)["weight"], would cover alpha. - Docs.
model-format-support.mdxlists 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.
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
left a comment
There was a problem hiding this comment.
Thanks, c818e18 resolves the blocking point. Rechecked on 189c516:
- Real file from #10018. Car Meet Garage through
LoRALoader._load_modelnow 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 1052LoKRLayers, 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_scaleand 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 istest_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_countasserts on the loader's mock logger, while the per-layer warning came from the module logger inlayers/utils.py. With that warning restored, this test would still pass, and thecaplogassertions 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_SUFFIXESincludes.magnitudeand.lora_magnitude_vector.weight, but_without_unsupported_magnitudeonly removesdora_scaleanddora_magnitude. Krea-2 renames both first, so it stays at one warning. On kohya keys they reachLoKRLayer/LoHALayerunrenamed, and each layer then logsUnexpected keys found. I only got there with synthetic files. LyCORIS writesdora_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.
Summary
Fixes #10018
Root cause
The shared layer factory dispatched any state dict containing
dora_scaleordora_magnitudetoDoRALayerbefore checking for LoKr or LoHa factors. DoKr and DoHa layers therefore tried to read nonexistentlora_up.weight/lora_down.weighttensors and failed withKeyError.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 -q121 passedcheckandformat --checkpass 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
What's Newcopy (if doing a release after this PR)Need help on this PR? Tag
@codesmith-botwith what you need.