Repository navigation
fix(lora): identify and convert FLUX kohya and OneTrainer layers outside the blocks - #10206
Pfannkuchensack wants to merge 2 commits into
Conversation
…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
| elif not re.fullmatch(FLUX_KOHYA_NON_BLOCK_KEY_REGEX, k): | ||
| return False | ||
| return anchored |
There was a problem hiding this comment.
[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.
There was a problem hiding this comment.
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
linear1gives the hidden size (input width) and the MLP ratio. FLUX.1's output is3072 * (3 + 4) = 21504, while Klein 4B's is 27648. The variant is read from the same widths when there is noimg_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_w1onlinear1, both widths.
There was a problem hiding this comment.
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.
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.img_in,txt_in,time_in,vector_in,guidance_in,final_layer_linear,final_layer_adaLN_modulation_1and 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 fornorm_out.linear. FLUX.2 [dev] also gains itsguidance_in→time_guidance_embed.guidance_embedderrename.x_embedder,context_embedder,norm_out, the finalproj_outand thetime_text_embedlayers are accepted next to the blocks. The existing diffusers converter handles them, with the half swap fornorm_out.linear.linear1width gives the hidden size and MLP ratio (FLUX.1: 21504 wide, Klein 4B: 27648). kohya FLUX.2 LoRAs without double blocks ortxt_inno longer fall back to FLUX.1.insert_periods_into_kohya_key: splitting a key now backtracks. Before, it tookff/norm1and then failed onff_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.FluxLoRASaverLEGACY output with the full layer filter, written from a meta model with random values.LastLayerupdate.main.Failed to find modulewarnings in the log.mainthe file does not install at all.Review
mainleaves them unidentified. This is fixed with the single-block requirement and covered by tests.single_stream_modulation_linwas 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 wholelokr_w1.main, OneTrainer legacy SD3.5/Qwen-Image files that train only blocks are identified as FLUX.1. This PR does not change that.ff_linear_in,to_qkv_mlp_proj, …) is still not supported.Checklist
What's Newcopy (if doing a release after this PR)Open in [code]smith · Need help on this PR? Tag
@codesmith-botwith what you need.