Skip to content

fix(lora): identify and convert FLUX kohya and OneTrainer layers outside the blocks - #10206

Open
Pfannkuchensack wants to merge 2 commits into
invoke-ai:mainfrom
Pfannkuchensack:fix/flux-lora-non-block-layers
Open

Pfannkuchensack wants to merge 2 commits into
invoke-ai:mainfrom
Pfannkuchensack:fix/flux-lora-non-block-layers

Conversation

@Pfannkuchensack

@Pfannkuchensack Pfannkuchensack commented Nov 8, 2026 •

Copy link
Copy Markdown
Member

Summary

FLUX kohya (lora_unet_*) and OneTrainer legacy (lora_transformer_*) LoRAs that also train layers outside the transformer blocks installed as Unknown, and forcing the base failed in the converters. Such files now load, and every layer patches the right module.

  • kohya: img_in, txt_in, time_in, vector_in, guidance_in, final_layer_linear, final_layer_adaLN_modulation_1 and the FLUX.2 shared modulations map to BFL module paths through one table. The same table drives the detector. The converter no longer drops the final layer. On FLUX.2 the existing BFL→diffusers rename applies them, including the half swap for norm_out.linear. FLUX.2 [dev] also gains its guidance_in → time_guidance_embed.guidance_embedder rename.
  • OneTrainer: x_embedder, context_embedder, norm_out, the final proj_out and the time_text_embed layers are accepted next to the blocks. The existing diffusers converter handles them, with the half swap for norm_out.linear.
  • Guard against other families: non-block layers alone never identify FLUX, because Qwen-Image and SD3 share these names. kohya needs a FLUX block, CLIP or T5 key. OneTrainer needs a single-block key, because SD3's joint blocks have the same names as FLUX's double blocks.
  • FLUX.2 identification (kohya): the shared stream modulations mark a file as FLUX.2, and a single block's fused linear1 width gives the hidden size and MLP ratio (FLUX.1: 21504 wide, Klein 4B: 27648). kohya FLUX.2 LoRAs without double blocks or txt_in no longer fall back to FLUX.1.
  • insert_periods_into_kohya_key: splitting a key now backtracks. Before, it took ff / norm1 and then failed on ff_context / norm1_context, so every OneTrainer FLUX.1 LoRA that trains those layers failed to load (OneTrainer's FLUX.1 preset trains all layers). Keys that parsed before still parse the same way.

Before → after, with the real PJMixers FLUX.2 Klein 4B distill extract (rank 16, kohya): Unknown → FLUX.2 Klein 4B, all 89 layers loaded, 109/109 diffusers-side patches resolved.

Related Issues / Discussions

Closes #10030

Follows #10052 (final-layer half swap) and #10045 (guidance_in).

QA Instructions

  • pytest tests/backend/patches tests/backend/model_manager/configs: 2148 passed. The new and changed test files also pass on their own. Ruff check and format are clean.
  • New fixtures:
    • The header of the real PJMixers file.
    • OneTrainer's own FluxLoRASaver LEGACY output with the full layer filter, written from a meta model with random values.
  • Tests cover:
    • Detection, conversion, and resolution on meta FLUX.1 and FLUX.2 models, including FLUX.2 [dev] with a guidance embedder.
    • Final-layer half order on FLUX.1 and FLUX.2, checked against a hand-computed BFL LastLayer update.
    • SD3, Qwen-Image and kohya-Qwen key sets staying non-FLUX, checked with the detectors and the LoRA config probes.
  • Identification with all LoRA config probes on OneTrainer-legacy key sets built from full-size diffusers SD3.5 Large/Medium and Qwen-Image: full-filter and blocks-plus-top-level sets stay unidentified, as on main.
  • GPU (RTX 4090), FLUX.2 Klein 4B Base with the PJMixers extract, 4 steps, CFG 1, 1 warm-up plus 3 seeds:
    • Without the LoRA the images stay hazy and blurred. With it they are sharp, like the distilled Klein 4B.
    • No Failed to find module warnings in the log.
    • On main the file does not install at all.
  • Not run: an image test of a trained FLUX.1 LoRA with non-block layers. FLUX.1 is covered by load/resolve tests only.

Review

  • A review found that the first version of the OneTrainer detector identified full-filter SD3.5 OneTrainer files as FLUX.1, where main leaves them unidentified. This is fixed with the single-block requirement and covered by tests.
  • Maintainer review: a Klein 4B kohya subset with only a single block and single_stream_modulation_lin was identified as FLUX.1 and failed at apply time. Fixed in the identification and covered with partial-layer tests from the real fixture, a FLUX.1 counter-check, and LoKr with a whole lokr_w1.
  • Remaining limits:
    • Already on main, OneTrainer legacy SD3.5/Qwen-Image files that train only blocks are identified as FLUX.1. This PR does not change that.
    • OneTrainer's legacy layout for FLUX.2 (ff_linear_in, to_qkv_mlp_proj, …) is still not supported.
    • In the PJMixers extract the final-layer layers change the image only slightly. The half order rests on the numeric tests, not on an image comparison.

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)

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

…ide the blocks

kohya and OneTrainer FLUX LoRAs that also train img_in/txt_in, the embedders, the final layer or the FLUX.2 modulations now load instead of installing as Unknown.
OneTrainer non-block layers need single blocks alongside, so SD3 and Qwen-Image files stay out; FLUX.2 [dev] gets its guidance embedder rename.
Kohya key splitting no longer picks `ff`/`norm1` over `ff_context`/`norm1_context`.

Closes invoke-ai#10030
@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 docs PRs that change docs labels Nov 8, 2026
@lstein lstein self-assigned this Nov 8, 2026
@lstein lstein added the 7.0.1 label Nov 8, 2026
@lstein lstein moved this to 7.0.0: UX overhaul in Invoke - Community Roadmap Nov 8, 2026
@lstein lstein added 7.0.0 and removed 7.0.1 labels Nov 8, 2026
Comment on lines +93 to +95
elif not re.fullmatch(FLUX_KOHYA_NON_BLOCK_KEY_REGEX, k):
return False
return anchored

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

[P2] Newly accepted FLUX.2 LoRAs can install as FLUX.1.

Accepts shared modulation layers, but family identification doesn’t recognize those FLUX.2-specific names. Reproduced with the PR’s Klein 4B fixture restricted to single_blocks_0_linear1, single_blocks_0_linear2, and single_stream_modulation_lin. Previously unrecognized, it now identifies as FLUX.1. Its patch then fails with a shape mismatch: (27648, 3072) versus (21504, 3072).

Extend family/variant identification alongside the format detector and add a partial-layer regression test.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Thanks, reproduced: that subset was claimed by the FLUX.1 config. Fixed in ea78cfb.

kohya FLUX.2 identification only looked at txt_in, vector_in and the double-block widths, so a file without those fell back to FLUX.1. It now also uses two signals:

  • The shared stream modulations (double_stream_modulation_*, single_stream_modulation_lin) exist only in FLUX.2.
  • A single block's fused linear1 gives the hidden size (input width) and the MLP ratio. FLUX.1's output is 3072 * (3 + 4) = 21504, while Klein 4B's is 27648. The variant is read from the same widths when there is no img_attn_proj. This also covers single-block-only kohya FLUX.2 files, which hit the same mismatch before this PR.

Regression tests in test_flux3_lora_probe_independence.py use subsets of the real Klein 4B fixture:

  • your exact set;
  • single block only;
  • modulation plus txt_attn;
  • a FLUX.1 single-block counter-check from the FLUX.1 kohya fixture;
  • LoKr with a whole lokr_w1 on linear1, both widths.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

The original shared-modulation reproducer is fixed, but other newly accepted FLUX.2 subsets still identify as FLUX.1.

Restrict the PR’s Klein 4B fixture to img_in and double_blocks_0_img_attn_proj. The expanded detector accepts it, but family identification ignores img_in’s distinguishing width and falls back to FLUX.1. Applying its (3072,128) input patch to FLUX.1’s (3072,64) module fails.

Extend family discrimination—or reject incompatible ambiguity—and cover this partial filter.

… and modulations

kohya FLUX.2 LoRAs without double blocks or txt_in were identified as FLUX.1 and failed to apply.
The shared stream modulations now mark FLUX.2, and a single block's fused linear1 width gives hidden size and MLP ratio.

This branch has not been deployed

No deployments
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 docs PRs that change docs 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]: FLUX kohya and OneTrainer LoRAs with layers outside the transformer blocks are not identified or converted

3 participants