Add Flux2 LoKR adapter support via peft and fix alpha handling

Support native LoKR loading for Flux2/Klein models through peft's
LoKrConfig + inject_adapter_in_model, matching how LoRA is already
handled. Also fix lora_down/lora_up + alpha keys being left unconsumed.

Changes:
- Fix alpha baking for lora_down/lora_up format in the Flux2 converter
- Add LoKR conversion function mapping BFL keys to diffusers paths
- Use fuse-first QKV strategy: fuse model projections so LoKR targets
  fused to_qkv/to_added_qkv directly (avoids lossy Kronecker splitting)
- Detect and infer decompose_factor from checkpoint w1 shapes
- Add LoKR detection branch in load_lora_adapter with LoKrConfig creation
- Add re-fusion helper for loading LoRA onto already-fused models
- Temporary debug logging with [LoKR DEBUG] prefix

GitHub issue: huggingface#13261

Author: CalamitousFelicitousness

src/diffusers/loaders/lora_conversion_utils.py

@@ -2331,6 +2331,18 @@ def _convert_non_diffusers_flux2_lora_to_diffusers(state_dict):
             temp_state_dict[new_key] = v
         original_state_dict = temp_state_dict
 
+        # Bake alpha/rank scaling into lora_A weights (matches Flux1 kohya alpha handling).
+        # Without this, .alpha keys remain unconsumed and trigger a validation error.
+        alpha_keys = [k for k in original_state_dict if k.endswith(".alpha")]
+        for alpha_key in alpha_keys:
+            alpha = original_state_dict.pop(alpha_key).item()
+            module_path = alpha_key[: -len(".alpha")]
+            lora_a_key = f"{module_path}.lora_A.weight"
+            if lora_a_key in original_state_dict:
+                rank = original_state_dict[lora_a_key].shape[0]
+                scale = alpha / rank
+                original_state_dict[lora_a_key] = original_state_dict[lora_a_key] * scale
+
     num_double_layers = 0
     num_single_layers = 0
     for key in original_state_dict.keys():
@@ -2443,6 +2455,166 @@ def _convert_non_diffusers_flux2_lora_to_diffusers(state_dict):
     return converted_state_dict
 
 
+def _convert_non_diffusers_flux2_lokr_to_diffusers(state_dict):
+    """Convert non-diffusers Flux2 LoKR state dict (kohya/LyCORIS format) to peft-compatible diffusers format.
+
+    Uses fuse-first QKV mapping: BFL fused `img_attn.qkv` maps to diffusers `attn.to_qkv` (created by
+    `fuse_projections()`), avoiding lossy Kronecker factor splitting. The caller must fuse the model's
+    QKV projections before injecting the adapter.
+    """
+    converted_state_dict = {}
+
+    prefix = "diffusion_model."
+    original_state_dict = {k[len(prefix) :] if k.startswith(prefix) else k: v for k, v in state_dict.items()}
+
+    # Log key patterns for debugging
+    suffixes = set()
+    for k in original_state_dict:
+        for s in ("lokr_w1_a", "lokr_w1_b", "lokr_w2_a", "lokr_w2_b", "lokr_w1", "lokr_w2", "lokr_t2", "alpha"):
+            if k.endswith(f".{s}"):
+                suffixes.add(s)
+                break
+    logger.warning(
+        f"[LoKR DEBUG] LoKR converter: {len(original_state_dict)} keys, weight types: {sorted(suffixes)}"
+    )
+
+    num_double_layers = 0
+    num_single_layers = 0
+    for key in original_state_dict:
+        if key.startswith("single_blocks."):
+            num_single_layers = max(num_single_layers, int(key.split(".")[1]) + 1)
+        elif key.startswith("double_blocks."):
+            num_double_layers = max(num_double_layers, int(key.split(".")[1]) + 1)
+
+    lokr_suffixes = ("lokr_w1", "lokr_w1_a", "lokr_w1_b", "lokr_w2", "lokr_w2_a", "lokr_w2_b", "lokr_t2")
+
+    def _remap_lokr_module(bfl_path, diff_path):
+        """Pop all lokr keys for a BFL module, bake alpha scaling, and store under diffusers path."""
+        alpha_key = f"{bfl_path}.alpha"
+        alpha = original_state_dict.pop(alpha_key).item() if alpha_key in original_state_dict else None
+
+        for suffix in lokr_suffixes:
+            src_key = f"{bfl_path}.{suffix}"
+            if src_key not in original_state_dict:
+                continue
+
+            weight = original_state_dict.pop(src_key)
+
+            # Bake alpha/rank scaling into the first w1 tensor encountered for this module.
+            # After baking, peft's config uses alpha=r so its runtime scaling is 1.0.
+            if alpha is not None and suffix in ("lokr_w1", "lokr_w1_a"):
+                w2a_key = f"{bfl_path}.lokr_w2_a"
+                w1a_key = f"{bfl_path}.lokr_w1_a"
+                if w2a_key in original_state_dict:
+                    r_eff = original_state_dict[w2a_key].shape[1]
+                elif w1a_key in original_state_dict:
+                    r_eff = original_state_dict[w1a_key].shape[1]
+                else:
+                    r_eff = alpha
+                scale = alpha / r_eff
+                logger.warning(
+                    f"[LoKR DEBUG] Alpha bake: {bfl_path} alpha={alpha:.1f} r={r_eff} scale={scale:.4f}"
+                )
+                weight = weight * scale
+                alpha = None  # only bake once per module
+
+            converted_state_dict[f"{diff_path}.{suffix}"] = weight
+
+    # --- Single blocks ---
+    for sl in range(num_single_layers):
+        _remap_lokr_module(f"single_blocks.{sl}.linear1", f"single_transformer_blocks.{sl}.attn.to_qkv_mlp_proj")
+        _remap_lokr_module(f"single_blocks.{sl}.linear2", f"single_transformer_blocks.{sl}.attn.to_out")
+
+    # --- Double blocks ---
+    for dl in range(num_double_layers):
+        tb = f"transformer_blocks.{dl}"
+        db = f"double_blocks.{dl}"
+
+        # QKV → fused to_qkv / to_added_qkv (model must be fused before injection)
+        _remap_lokr_module(f"{db}.img_attn.qkv", f"{tb}.attn.to_qkv")
+        _remap_lokr_module(f"{db}.txt_attn.qkv", f"{tb}.attn.to_added_qkv")
+
+        # Projections
+        _remap_lokr_module(f"{db}.img_attn.proj", f"{tb}.attn.to_out.0")
+        _remap_lokr_module(f"{db}.txt_attn.proj", f"{tb}.attn.to_add_out")
+
+        # MLPs
+        _remap_lokr_module(f"{db}.img_mlp.0", f"{tb}.ff.linear_in")
+        _remap_lokr_module(f"{db}.img_mlp.2", f"{tb}.ff.linear_out")
+        _remap_lokr_module(f"{db}.txt_mlp.0", f"{tb}.ff_context.linear_in")
+        _remap_lokr_module(f"{db}.txt_mlp.2", f"{tb}.ff_context.linear_out")
+
+    # --- Extra mappings (embedders, modulation, final layer) ---
+    extra_mappings = {
+        "img_in": "x_embedder",
+        "txt_in": "context_embedder",
+        "time_in.in_layer": "time_guidance_embed.timestep_embedder.linear_1",
+        "time_in.out_layer": "time_guidance_embed.timestep_embedder.linear_2",
+        "final_layer.linear": "proj_out",
+        "final_layer.adaLN_modulation.1": "norm_out.linear",
+        "single_stream_modulation.lin": "single_stream_modulation.linear",
+        "double_stream_modulation_img.lin": "double_stream_modulation_img.linear",
+        "double_stream_modulation_txt.lin": "double_stream_modulation_txt.linear",
+    }
+    for bfl_key, diff_key in extra_mappings.items():
+        _remap_lokr_module(bfl_key, diff_key)
+
+    if len(original_state_dict) > 0:
+        raise ValueError(f"`original_state_dict` should be empty at this point but has {original_state_dict.keys()=}.")
+
+    for key in list(converted_state_dict.keys()):
+        converted_state_dict[f"transformer.{key}"] = converted_state_dict.pop(key)
+
+    return converted_state_dict
+
+
+def _refuse_flux2_lora_state_dict(state_dict):
+    """Re-fuse separate Q/K/V LoRA keys into fused to_qkv/to_added_qkv keys.
+
+    When the model's QKV projections are fused (e.g., from a prior LoKR load), incoming LoRA keys
+    that target separate to_q/to_k/to_v must be re-fused to match. This is the exact inverse of the
+    QKV split in ``_convert_non_diffusers_flux2_lora_to_diffusers``.
+    """
+    converted = {}
+    remaining = dict(state_dict)
+
+    # Detect double block indices from keys
+    num_double_layers = 0
+    for key in remaining:
+        if ".transformer_blocks." in key:
+            parts = key.split(".")
+            idx = parts.index("transformer_blocks") + 1
+            if idx < len(parts) and parts[idx].isdigit():
+                num_double_layers = max(num_double_layers, int(parts[idx]) + 1)
+
+    # Fuse Q/K/V for image stream and text stream
+    qkv_groups = [
+        (["to_q", "to_k", "to_v"], "to_qkv"),
+        (["add_q_proj", "add_k_proj", "add_v_proj"], "to_added_qkv"),
+    ]
+
+    for dl in range(num_double_layers):
+        attn_prefix = f"transformer.transformer_blocks.{dl}.attn"
+        for separate_keys, fused_name in qkv_groups:
+            for lora_key in ("lora_A", "lora_B"):
+                src_keys = [f"{attn_prefix}.{sk}.{lora_key}.weight" for sk in separate_keys]
+                if not all(k in remaining for k in src_keys):
+                    continue
+
+                weights = [remaining.pop(k) for k in src_keys]
+                dst_key = f"{attn_prefix}.{fused_name}.{lora_key}.weight"
+                if lora_key == "lora_A":
+                    # lora_A was replicated during split — all three are identical, take the first
+                    converted[dst_key] = weights[0]
+                else:
+                    # lora_B was chunked along dim=0 — concatenate back
+                    converted[dst_key] = torch.cat(weights, dim=0)
+
+    # Pass through all non-QKV keys unchanged
+    converted.update(remaining)
+    return converted
+
+
 def _convert_non_diffusers_z_image_lora_to_diffusers(state_dict):
     """
     Convert non-diffusers ZImage LoRA state dict to diffusers format.

src/diffusers/loaders/lora_pipeline.py

@@ -45,6 +45,7 @@
     _convert_hunyuan_video_lora_to_diffusers,
     _convert_kohya_flux_lora_to_diffusers,
     _convert_musubi_wan_lora_to_diffusers,
+    _convert_non_diffusers_flux2_lokr_to_diffusers,
     _convert_non_diffusers_flux2_lora_to_diffusers,
     _convert_non_diffusers_hidream_lora_to_diffusers,
     _convert_non_diffusers_lora_to_diffusers,
@@ -56,6 +57,7 @@
     _convert_non_diffusers_z_image_lora_to_diffusers,
     _convert_xlabs_flux_lora_to_diffusers,
     _maybe_map_sgm_blocks_to_diffusers,
+    _refuse_flux2_lora_state_dict,
 )
 
 
@@ -5679,12 +5681,26 @@ def lora_state_dict(
 
         is_ai_toolkit = any(k.startswith("diffusion_model.") for k in state_dict)
         if is_ai_toolkit:
-            state_dict = _convert_non_diffusers_flux2_lora_to_diffusers(state_dict)
+            is_lokr = any("lokr_" in k for k in state_dict)
+            if is_lokr:
+                logger.warning(f"[LoKR DEBUG] Detected LoKR format ({len(state_dict)} keys), converting to diffusers")
+                state_dict = _convert_non_diffusers_flux2_lokr_to_diffusers(state_dict)
+                logger.warning(f"[LoKR DEBUG] Converted to {len(state_dict)} diffusers keys")
+                if metadata is None:
+                    metadata = {}
+                metadata["is_lokr"] = "true"
+            else:
+                has_alpha = any(k.endswith(".alpha") for k in state_dict)
+                has_down_up = any("lora_down" in k or "lora_up" in k for k in state_dict)
+                logger.warning(
+                    f"[LoKR DEBUG] Detected LoRA format ({len(state_dict)} keys, "
+                    f"has_alpha={has_alpha}, has_lora_down_up={has_down_up})"
+                )
+                state_dict = _convert_non_diffusers_flux2_lora_to_diffusers(state_dict)
 
         out = (state_dict, metadata) if return_lora_metadata else state_dict
         return out
 
-    # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights
     def load_lora_weights(
         self,
         pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
@@ -5712,13 +5728,29 @@ def load_lora_weights(
         kwargs["return_lora_metadata"] = True
         state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs)
 
-        is_correct_format = all("lora" in key for key in state_dict.keys())
+        is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys())
         if not is_correct_format:
-            raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.")
+            raise ValueError("Invalid LoRA/LoKR checkpoint. Make sure all param names contain `'lora'` or `'lokr'`.")
+
+        # For LoKR adapters, fuse QKV projections so peft can target the fused modules directly.
+        is_lokr = metadata is not None and metadata.get("is_lokr") == "true"
+        transformer = getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer
+        if is_lokr:
+            logger.warning("[LoKR DEBUG] Fusing QKV projections for LoKR adapter")
+            transformer.fuse_qkv_projections()
+        elif (
+            hasattr(transformer, "transformer_blocks")
+            and len(transformer.transformer_blocks) > 0
+            and getattr(transformer.transformer_blocks[0].attn, "fused_projections", False)
+        ):
+            # Model is fused (e.g., from a prior LoKR load) but this is a LoRA adapter
+            # targeting separate Q/K/V. Re-fuse the LoRA keys to match the fused modules.
+            logger.warning("[LoKR DEBUG] Model is fused, re-fusing LoRA keys to match fused QKV modules")
+            state_dict = _refuse_flux2_lora_state_dict(state_dict)
 
         self.load_lora_into_transformer(
             state_dict,
-            transformer=getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer,
+            transformer=transformer,
             adapter_name=adapter_name,
             metadata=metadata,
             _pipeline=self,

src/diffusers/loaders/peft.py

@@ -37,7 +37,7 @@
     set_adapter_layers,
     set_weights_and_activate_adapters,
 )
-from ..utils.peft_utils import _create_lora_config, _maybe_warn_for_unhandled_keys
+from ..utils.peft_utils import _create_lokr_config, _create_lora_config, _maybe_warn_for_unhandled_keys
 from .lora_base import _fetch_state_dict, _func_optionally_disable_offloading
 from .unet_loader_utils import _maybe_expand_lora_scales
 
@@ -232,56 +232,73 @@ def load_lora_adapter(
                     "Please choose an existing adapter name or set `hotswap=False` to prevent hotswapping."
                 )
 
-            # check with first key if is not in peft format
-            first_key = next(iter(state_dict.keys()))
-            if "lora_A" not in first_key:
-                state_dict = convert_unet_state_dict_to_peft(state_dict)
-
-            # Control LoRA from SAI is different from BFL Control LoRA
-            # https://huggingface.co/stabilityai/control-lora
-            # https://huggingface.co/comfyanonymous/ControlNet-v1-1_fp16_safetensors
-            is_sai_sd_control_lora = "lora_controlnet" in state_dict
-            if is_sai_sd_control_lora:
-                state_dict = convert_sai_sd_control_lora_state_dict_to_peft(state_dict)
-
-            rank = {}
-            for key, val in state_dict.items():
-                # Cannot figure out rank from lora layers that don't have at least 2 dimensions.
-                # Bias layers in LoRA only have a single dimension
-                if "lora_B" in key and val.ndim > 1:
-                    # Check out https://github.com/huggingface/peft/pull/2419 for the `^` symbol.
-                    # We may run into some ambiguous configuration values when a model has module
-                    # names, sharing a common prefix (`proj_out.weight` and `blocks.transformer.proj_out.weight`,
-                    # for example) and they have different LoRA ranks.
-                    rank[f"^{key}"] = val.shape[1]
-
-            if network_alphas is not None and len(network_alphas) >= 1:
-                alpha_keys = [k for k in network_alphas.keys() if k.startswith(f"{prefix}.")]
-                network_alphas = {
-                    k.removeprefix(f"{prefix}."): v for k, v in network_alphas.items() if k in alpha_keys
-                }
-
-            # adapter_name
-            if adapter_name is None:
-                adapter_name = get_adapter_name(self)
-
-            # create LoraConfig
-            lora_config = _create_lora_config(
-                state_dict,
-                network_alphas,
-                metadata,
-                rank,
-                model_state_dict=self.state_dict(),
-                adapter_name=adapter_name,
-            )
+            # Detect whether this is a LoKR adapter (Kronecker product, not low-rank)
+            is_lokr = any("lokr_" in k for k in state_dict)
+
+            if is_lokr:
+                logger.warning(f"[LoKR DEBUG] load_lora_adapter: detected LoKR state dict ({len(state_dict)} keys)")
+                if adapter_name is None:
+                    adapter_name = get_adapter_name(self)
+                lora_config = _create_lokr_config(state_dict)
+                logger.warning(
+                    f"[LoKR DEBUG] LoKrConfig: r={lora_config.r}, alpha={lora_config.alpha}, "
+                    f"decompose_both={lora_config.decompose_both}, "
+                    f"decompose_factor={lora_config.decompose_factor}, "
+                    f"targets={len(lora_config.target_modules)} modules, "
+                    f"rank_pattern={dict(lora_config.rank_pattern) if lora_config.rank_pattern else '{}'}"
+                )
+                is_sai_sd_control_lora = False
+            else:
+                # check with first key if is not in peft format
+                first_key = next(iter(state_dict.keys()))
+                if "lora_A" not in first_key:
+                    state_dict = convert_unet_state_dict_to_peft(state_dict)
+
+                # Control LoRA from SAI is different from BFL Control LoRA
+                # https://huggingface.co/stabilityai/control-lora
+                # https://huggingface.co/comfyanonymous/ControlNet-v1-1_fp16_safetensors
+                is_sai_sd_control_lora = "lora_controlnet" in state_dict
+                if is_sai_sd_control_lora:
+                    state_dict = convert_sai_sd_control_lora_state_dict_to_peft(state_dict)
+
+                rank = {}
+                for key, val in state_dict.items():
+                    # Cannot figure out rank from lora layers that don't have at least 2 dimensions.
+                    # Bias layers in LoRA only have a single dimension
+                    if "lora_B" in key and val.ndim > 1:
+                        # Check out https://github.com/huggingface/peft/pull/2419 for the `^` symbol.
+                        # We may run into some ambiguous configuration values when a model has module
+                        # names, sharing a common prefix (`proj_out.weight` and `blocks.transformer.proj_out.weight`,
+                        # for example) and they have different LoRA ranks.
+                        rank[f"^{key}"] = val.shape[1]
+
+                if network_alphas is not None and len(network_alphas) >= 1:
+                    alpha_keys = [k for k in network_alphas.keys() if k.startswith(f"{prefix}.")]
+                    network_alphas = {
+                        k.removeprefix(f"{prefix}."): v for k, v in network_alphas.items() if k in alpha_keys
+                    }
+
+                # adapter_name
+                if adapter_name is None:
+                    adapter_name = get_adapter_name(self)
+
+                # create LoraConfig
+                lora_config = _create_lora_config(
+                    state_dict,
+                    network_alphas,
+                    metadata,
+                    rank,
+                    model_state_dict=self.state_dict(),
+                    adapter_name=adapter_name,
+                )
 
-            # Adjust LoRA config for Control LoRA
-            if is_sai_sd_control_lora:
-                lora_config.lora_alpha = lora_config.r
-                lora_config.alpha_pattern = lora_config.rank_pattern
-                lora_config.bias = "all"
-                lora_config.modules_to_save = lora_config.exclude_modules
-                lora_config.exclude_modules = None
+                # Adjust LoRA config for Control LoRA
+                if is_sai_sd_control_lora:
+                    lora_config.lora_alpha = lora_config.r
+                    lora_config.alpha_pattern = lora_config.rank_pattern
+                    lora_config.bias = "all"
+                    lora_config.modules_to_save = lora_config.exclude_modules
+                    lora_config.exclude_modules = None
 
             # <Unsafe code
             # We can be sure that the following works as it just sets attention processors, lora layers and puts all in the same dtype