From 00717303410b65948d5431911807837be4d81d25 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Mon, 23 Mar 2026 04:47:13 +0100 Subject: [PATCH 1/2] feat: add support for OneTrainer BFL Flux LoRA format Newer versions of OneTrainer export Flux LoRAs using BFL internal key names (double_blocks, single_blocks, img_attn, etc.) with a 'transformer.' prefix and split QKV projections (qkv.0/1/2, linear1.0/1/2/3). This format was not recognized by any existing detector. Add detection and conversion for this format, merging split QKV and linear1 layers into MergedLayerPatch instances for the fused BFL model. --- .../model_manager/load/model_loaders/lora.py | 6 + invokeai/backend/model_manager/taxonomy.py | 1 + ...ux_onetrainer_bfl_lora_conversion_utils.py | 170 ++++++++++++++++++ .../patches/lora_conversions/formats.py | 5 + 4 files changed, 182 insertions(+) create mode 100644 invokeai/backend/patches/lora_conversions/flux_onetrainer_bfl_lora_conversion_utils.py diff --git a/invokeai/backend/model_manager/load/model_loaders/lora.py b/invokeai/backend/model_manager/load/model_loaders/lora.py index d39982456af..67d862a01d5 100644 --- a/invokeai/backend/model_manager/load/model_loaders/lora.py +++ b/invokeai/backend/model_manager/load/model_loaders/lora.py @@ -44,6 +44,10 @@ is_state_dict_likely_in_flux_kohya_format, lora_model_from_flux_kohya_state_dict, ) +from invokeai.backend.patches.lora_conversions.flux_onetrainer_bfl_lora_conversion_utils import ( + is_state_dict_likely_in_flux_onetrainer_bfl_format, + lora_model_from_flux_onetrainer_bfl_state_dict, +) from invokeai.backend.patches.lora_conversions.flux_onetrainer_lora_conversion_utils import ( is_state_dict_likely_in_flux_onetrainer_format, lora_model_from_flux_onetrainer_state_dict, @@ -128,6 +132,8 @@ def _load_model( model = lora_model_from_flux_diffusers_state_dict(state_dict=state_dict, alpha=None) elif is_state_dict_likely_in_flux_kohya_format(state_dict=state_dict): model = lora_model_from_flux_kohya_state_dict(state_dict=state_dict) + elif is_state_dict_likely_in_flux_onetrainer_bfl_format(state_dict=state_dict): + model = lora_model_from_flux_onetrainer_bfl_state_dict(state_dict=state_dict) elif is_state_dict_likely_in_flux_onetrainer_format(state_dict=state_dict): model = lora_model_from_flux_onetrainer_state_dict(state_dict=state_dict) elif is_state_dict_likely_flux_control(state_dict=state_dict): diff --git a/invokeai/backend/model_manager/taxonomy.py b/invokeai/backend/model_manager/taxonomy.py index c002418a6bd..9dc0da77335 100644 --- a/invokeai/backend/model_manager/taxonomy.py +++ b/invokeai/backend/model_manager/taxonomy.py @@ -210,6 +210,7 @@ class FluxLoRAFormat(str, Enum): AIToolkit = "flux.aitoolkit" XLabs = "flux.xlabs" BflPeft = "flux.bfl_peft" + OneTrainerBfl = "flux.onetrainer_bfl" AnyVariant: TypeAlias = Union[ diff --git a/invokeai/backend/patches/lora_conversions/flux_onetrainer_bfl_lora_conversion_utils.py b/invokeai/backend/patches/lora_conversions/flux_onetrainer_bfl_lora_conversion_utils.py new file mode 100644 index 00000000000..949af627dc8 --- /dev/null +++ b/invokeai/backend/patches/lora_conversions/flux_onetrainer_bfl_lora_conversion_utils.py @@ -0,0 +1,170 @@ +"""Utilities for detecting and converting FLUX LoRAs in OneTrainer BFL format. + +This format is produced by newer versions of OneTrainer and uses BFL internal key names +(double_blocks, single_blocks, img_attn, etc.) with a 'transformer.' prefix and +InvokeAI-native LoRA suffixes (lora_down.weight, lora_up.weight, alpha). + +Unlike the standard BFL PEFT format (which uses 'diffusion_model.' prefix and lora_A/lora_B), +this format also has split QKV projections: + - double_blocks.{i}.img_attn.qkv.{0,1,2} (Q, K, V separate) + - double_blocks.{i}.txt_attn.qkv.{0,1,2} (Q, K, V separate) + - single_blocks.{i}.linear1.{0,1,2,3} (Q, K, V, MLP separate) + +Example keys: + transformer.double_blocks.0.img_attn.qkv.0.lora_down.weight + transformer.double_blocks.0.img_attn.qkv.0.lora_up.weight + transformer.double_blocks.0.img_attn.qkv.0.alpha + transformer.single_blocks.0.linear1.3.lora_down.weight + transformer.double_blocks.0.img_mlp.0.lora_down.weight +""" + +import re +from typing import Any, Dict + +import torch + +from invokeai.backend.patches.layers.base_layer_patch import BaseLayerPatch +from invokeai.backend.patches.layers.merged_layer_patch import MergedLayerPatch, Range +from invokeai.backend.patches.layers.utils import any_lora_layer_from_state_dict +from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_TRANSFORMER_PREFIX +from invokeai.backend.patches.model_patch_raw import ModelPatchRaw + +_TRANSFORMER_PREFIX = "transformer." + +# Valid LoRA weight suffixes in this format. +_LORA_SUFFIXES = ("lora_down.weight", "lora_up.weight", "alpha") + +# Regex to detect split QKV keys in double blocks: e.g. "double_blocks.0.img_attn.qkv.1" +_SPLIT_QKV_RE = re.compile(r"^(double_blocks\.\d+\.(img_attn|txt_attn)\.qkv)\.\d+$") + +# Regex to detect split linear1 keys in single blocks: e.g. "single_blocks.0.linear1.2" +_SPLIT_LINEAR1_RE = re.compile(r"^(single_blocks\.\d+\.linear1)\.\d+$") + + +def is_state_dict_likely_in_flux_onetrainer_bfl_format( + state_dict: dict[str | int, Any], + metadata: dict[str, Any] | None = None, +) -> bool: + """Checks if the provided state dict is likely in the OneTrainer BFL FLUX LoRA format. + + This format uses BFL internal key names with 'transformer.' prefix and split QKV projections. + """ + str_keys = [k for k in state_dict.keys() if isinstance(k, str)] + if not str_keys: + return False + + # All keys must start with 'transformer.' + if not all(k.startswith(_TRANSFORMER_PREFIX) for k in str_keys): + return False + + # All keys must end with recognized LoRA suffixes. + if not all(k.endswith(_LORA_SUFFIXES) for k in str_keys): + return False + + # Must have BFL block structure (double_blocks or single_blocks) under transformer prefix. + has_bfl_blocks = any( + k.startswith("transformer.double_blocks.") or k.startswith("transformer.single_blocks.") for k in str_keys + ) + if not has_bfl_blocks: + return False + + # Must have split QKV pattern (qkv.0, qkv.1, qkv.2) to distinguish from other formats + # that might use transformer. prefix in the future. + has_split_qkv = any( + ".qkv.0." in k or ".qkv.1." in k or ".qkv.2." in k or ".linear1.0." in k for k in str_keys + ) + if not has_split_qkv: + return False + + return True + + +def _split_key(key: str) -> tuple[str, str]: + """Split a key into (layer_name, weight_suffix). + + Handles: + - 2-component suffixes ending with '.weight': e.g., 'lora_down.weight' → split at 2nd-to-last dot + - 1-component suffixes: e.g., 'alpha' → split at last dot + """ + if key.endswith(".weight"): + parts = key.rsplit(".", maxsplit=2) + return parts[0], f"{parts[1]}.{parts[2]}" + else: + parts = key.rsplit(".", maxsplit=1) + return parts[0], parts[1] + + +def lora_model_from_flux_onetrainer_bfl_state_dict(state_dict: Dict[str, torch.Tensor]) -> ModelPatchRaw: + """Convert a OneTrainer BFL format FLUX LoRA state dict to a ModelPatchRaw. + + Strips the 'transformer.' prefix, groups by layer, and merges split QKV/linear1 + layers into MergedLayerPatch instances. + """ + # Step 1: Strip prefix and group by layer name. + grouped_state_dict: dict[str, dict[str, torch.Tensor]] = {} + for key, value in state_dict.items(): + if not isinstance(key, str): + continue + + # Strip 'transformer.' prefix. + key = key[len(_TRANSFORMER_PREFIX) :] + + layer_name, suffix = _split_key(key) + + if layer_name not in grouped_state_dict: + grouped_state_dict[layer_name] = {} + grouped_state_dict[layer_name][suffix] = value + + # Step 2: Build LoRA layers, merging split QKV and linear1. + layers: dict[str, BaseLayerPatch] = {} + + # Identify which layers need merging. + merge_groups: dict[str, list[str]] = {} + standalone_keys: list[str] = [] + + for layer_key in grouped_state_dict: + qkv_match = _SPLIT_QKV_RE.match(layer_key) + linear1_match = _SPLIT_LINEAR1_RE.match(layer_key) + + if qkv_match: + parent = qkv_match.group(1) + if parent not in merge_groups: + merge_groups[parent] = [] + merge_groups[parent].append(layer_key) + elif linear1_match: + parent = linear1_match.group(1) + if parent not in merge_groups: + merge_groups[parent] = [] + merge_groups[parent].append(layer_key) + else: + standalone_keys.append(layer_key) + + # Process standalone layers. + for layer_key in standalone_keys: + layer_sd = grouped_state_dict[layer_key] + layers[f"{FLUX_LORA_TRANSFORMER_PREFIX}{layer_key}"] = any_lora_layer_from_state_dict(layer_sd) + + # Process merged layers. + for parent_key, sub_keys in merge_groups.items(): + # Sort by the numeric index at the end (e.g., qkv.0, qkv.1, qkv.2). + sub_keys.sort(key=lambda k: int(k.rsplit(".", maxsplit=1)[1])) + + sub_layers: list[BaseLayerPatch] = [] + sub_ranges: list[Range] = [] + dim_0_offset = 0 + + for sub_key in sub_keys: + layer_sd = grouped_state_dict[sub_key] + sub_layer = any_lora_layer_from_state_dict(layer_sd) + + # Determine the output dimension from the up weight shape. + up_weight = layer_sd["lora_up.weight"] + out_dim = up_weight.shape[0] + + sub_layers.append(sub_layer) + sub_ranges.append(Range(dim_0_offset, dim_0_offset + out_dim)) + dim_0_offset += out_dim + + layers[f"{FLUX_LORA_TRANSFORMER_PREFIX}{parent_key}"] = MergedLayerPatch(sub_layers, sub_ranges) + + return ModelPatchRaw(layers=layers) diff --git a/invokeai/backend/patches/lora_conversions/formats.py b/invokeai/backend/patches/lora_conversions/formats.py index 0b316602fcd..b3e00c288bd 100644 --- a/invokeai/backend/patches/lora_conversions/formats.py +++ b/invokeai/backend/patches/lora_conversions/formats.py @@ -14,6 +14,9 @@ from invokeai.backend.patches.lora_conversions.flux_kohya_lora_conversion_utils import ( is_state_dict_likely_in_flux_kohya_format, ) +from invokeai.backend.patches.lora_conversions.flux_onetrainer_bfl_lora_conversion_utils import ( + is_state_dict_likely_in_flux_onetrainer_bfl_format, +) from invokeai.backend.patches.lora_conversions.flux_onetrainer_lora_conversion_utils import ( is_state_dict_likely_in_flux_onetrainer_format, ) @@ -28,6 +31,8 @@ def flux_format_from_state_dict( ) -> FluxLoRAFormat | None: if is_state_dict_likely_in_flux_kohya_format(state_dict): return FluxLoRAFormat.Kohya + elif is_state_dict_likely_in_flux_onetrainer_bfl_format(state_dict, metadata): + return FluxLoRAFormat.OneTrainerBfl elif is_state_dict_likely_in_flux_onetrainer_format(state_dict): return FluxLoRAFormat.OneTrainer elif is_state_dict_likely_in_flux_diffusers_format(state_dict): From 119b2ea7cd5e3aaba7e42d0b3fdfe0a31a26fb05 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Mon, 23 Mar 2026 04:50:19 +0100 Subject: [PATCH 2/2] chore ruff --- .../flux_onetrainer_bfl_lora_conversion_utils.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/invokeai/backend/patches/lora_conversions/flux_onetrainer_bfl_lora_conversion_utils.py b/invokeai/backend/patches/lora_conversions/flux_onetrainer_bfl_lora_conversion_utils.py index 949af627dc8..b2109222a31 100644 --- a/invokeai/backend/patches/lora_conversions/flux_onetrainer_bfl_lora_conversion_utils.py +++ b/invokeai/backend/patches/lora_conversions/flux_onetrainer_bfl_lora_conversion_utils.py @@ -70,9 +70,7 @@ def is_state_dict_likely_in_flux_onetrainer_bfl_format( # Must have split QKV pattern (qkv.0, qkv.1, qkv.2) to distinguish from other formats # that might use transformer. prefix in the future. - has_split_qkv = any( - ".qkv.0." in k or ".qkv.1." in k or ".qkv.2." in k or ".linear1.0." in k for k in str_keys - ) + has_split_qkv = any(".qkv.0." in k or ".qkv.1." in k or ".qkv.2." in k or ".linear1.0." in k for k in str_keys) if not has_split_qkv: return False