From d58521b59a5543ddf129527a3e47fa22774bd400 Mon Sep 17 00:00:00 2001 From: Akhil Gupta Date: Tue, 6 Oct 2026 10:08:08 -0700 Subject: [PATCH 1/5] fix: derive text-encoder alpha_pattern keys from the LoRA module name --- src/diffusers/utils/peft_utils.py | 30 +++++++++++++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/src/diffusers/utils/peft_utils.py b/src/diffusers/utils/peft_utils.py index ea6f86798100..31ba150ac7ff 100644 --- a/src/diffusers/utils/peft_utils.py +++ b/src/diffusers/utils/peft_utils.py @@ -150,6 +150,34 @@ def unscale_lora_layers(model, weight: float | None = None): module.set_scale(adapter_name, 1.0) +# Maps the LoRA weight suffix of a text-encoder network alpha key (diffusers naming) +# to the corresponding PEFT module name suffix. PEFT's `alpha_pattern` keys must +# match the end of the PEFT module names (see `peft.utils.get_pattern_key`), e.g. +# the alpha key `text_model.encoder.layers.1.self_attn.to_k_lora.down.weight.alpha` +# belongs to the module `...self_attn.k_proj` -- not its parent block. +_TEXT_ENCODER_ALPHA_SUFFIX_TO_MODULE_SUFFIX = ( + (".to_q_lora.down.weight.alpha", ".q_proj"), + (".to_k_lora.down.weight.alpha", ".k_proj"), + (".to_v_lora.down.weight.alpha", ".v_proj"), + (".to_out_lora.down.weight.alpha", ".out_proj"), + (".lora_linear_layer.down.weight.alpha", ""), + # e.g. `text_model.text_projection.alpha` + (".alpha", ""), +) + + +def _text_encoder_alpha_pattern_key(alpha_key: str) -> str: + """Derive the PEFT `alpha_pattern` key from a text-encoder network alpha key.""" + for suffix, module_suffix in _TEXT_ENCODER_ALPHA_SUFFIX_TO_MODULE_SUFFIX: + if alpha_key.endswith(suffix): + key = alpha_key[: -len(suffix)] + module_suffix + # The loader strips the `text_model.` prefix for flattened text encoders, + # so the pattern key must not require it: `get_pattern_key` matches at the + # end of the module name, making the shorter key match either way. + return key.removeprefix("text_model.") + return alpha_key + + def get_peft_kwargs( rank_dict, network_alpha_dict, peft_state_dict, is_unet=True, model_state_dict=None, adapter_name=None ): @@ -186,7 +214,7 @@ def get_peft_kwargs( for k, v in alpha_pattern.items() } else: - alpha_pattern = {".".join(k.split(".down.")[0].split(".")[:-1]): v for k, v in alpha_pattern.items()} + alpha_pattern = {_text_encoder_alpha_pattern_key(k): v for k, v in alpha_pattern.items()} else: lora_alpha = set(network_alpha_dict.values()).pop() From e1e02ab069fe85ffc5a84f2dca2c5c4000c25317 Mon Sep 17 00:00:00 2001 From: Akhil Gupta Date: Tue, 6 Oct 2026 10:09:22 -0700 Subject: [PATCH 2/5] fix: derive text-encoder alpha_pattern keys from the LoRA module name --- tests/lora/test_peft_utils.py | 35 +++++++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/tests/lora/test_peft_utils.py b/tests/lora/test_peft_utils.py index 67b7e2695233..262c4fa41f93 100644 --- a/tests/lora/test_peft_utils.py +++ b/tests/lora/test_peft_utils.py @@ -77,3 +77,38 @@ def test_mixed_ranks_with_per_module_alphas_unchanged(): kwargs = get_peft_kwargs(_rank_dict(module_ranks), network_alphas, _peft_state_dict(module_ranks)) for module in module_ranks: assert _effective_scale(kwargs, module) == 1.0 + + +def test_text_encoder_alpha_pattern_names_lora_modules(): + # Regression test for https://github.com/huggingface/diffusers/issues/14970: + # text-encoder `alpha_pattern` keys must name the LoRA module itself (e.g. `k_proj`), + # not its parent block, otherwise PEFT's end-of-name matching misses and the + # default (most common) alpha is applied to every module. + peft_state_dict = { + "text_model.encoder.layers.0.self_attn.q_proj.lora_A.weight": None, + "text_model.encoder.layers.0.self_attn.q_proj.lora_B.weight": None, + "text_model.encoder.layers.1.self_attn.k_proj.lora_A.weight": None, + "text_model.encoder.layers.1.self_attn.k_proj.lora_B.weight": None, + } + rank_dict = {k.replace(".lora_A.", ".lora_B."): 4 for k in peft_state_dict if ".lora_A." in k} + network_alphas = { + "text_model.encoder.layers.0.self_attn.to_q_lora.down.weight.alpha": 4.0, + "text_model.encoder.layers.1.self_attn.to_k_lora.down.weight.alpha": 8.0, + } + kwargs = get_peft_kwargs(rank_dict, network_alphas, peft_state_dict, is_unet=False) + assert kwargs["lora_alpha"] == 4.0 + assert kwargs["alpha_pattern"] == {"encoder.layers.1.self_attn.k_proj": 8.0} + assert _effective_scale(kwargs, "encoder.layers.1.self_attn.k_proj") == 2.0 + assert _effective_scale(kwargs, "encoder.layers.0.self_attn.q_proj") == 1.0 + + +def test_text_encoder_alpha_pattern_handles_mlp_and_text_projection(): + peft_state_dict = {"m.lora_A.weight": None, "m.lora_B.weight": None} + network_alphas = { + "text_model.encoder.layers.0.mlp.fc1.lora_linear_layer.down.weight.alpha": 8.0, + "text_model.text_projection.alpha": 8.0, + "text_model.encoder.layers.0.mlp.fc2.lora_linear_layer.down.weight.alpha": 4.0, + } + kwargs = get_peft_kwargs({"m.lora_B.weight": 4}, network_alphas, peft_state_dict, is_unet=False) + assert kwargs["lora_alpha"] == 8.0 + assert kwargs["alpha_pattern"] == {"encoder.layers.0.mlp.fc2": 4.0} From cae69d9206064cba48342a022c1ab42a219fcaec Mon Sep 17 00:00:00 2001 From: Akhil Gupta Date: Mon, 12 Oct 2026 21:11:32 -0700 Subject: [PATCH 3/5] upload updated peft_utils.py --- .../utils/diffusers14972-peft-utils.py | 566 ++++++++++++++++++ 1 file changed, 566 insertions(+) create mode 100644 src/diffusers/utils/diffusers14972-peft-utils.py diff --git a/src/diffusers/utils/diffusers14972-peft-utils.py b/src/diffusers/utils/diffusers14972-peft-utils.py new file mode 100644 index 000000000000..9c2b3e40efdd --- /dev/null +++ b/src/diffusers/utils/diffusers14972-peft-utils.py @@ -0,0 +1,566 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +PEFT utilities: Utilities related to peft library +""" + +import collections +import functools +import importlib + +from packaging import version + +from . import logging +from .import_utils import is_peft_available, is_torch_available +from .torch_utils import empty_device_cache + + +logger = logging.get_logger(__name__) + +if is_torch_available(): + import torch + + +def recurse_remove_peft_layers(model): + r""" + Recursively replace all instances of `LoraLayer` with corresponding new layers in `model`. + """ + from peft.tuners.tuners_utils import BaseTunerLayer + + has_base_layer_pattern = False + for module in model.modules(): + if isinstance(module, BaseTunerLayer): + has_base_layer_pattern = hasattr(module, "base_layer") + break + + if has_base_layer_pattern: + from peft.utils import _get_submodules + + key_list = [key for key, _ in model.named_modules() if "lora" not in key] + for key in key_list: + try: + parent, target, target_name = _get_submodules(model, key) + except AttributeError: + continue + if hasattr(target, "base_layer"): + setattr(parent, target_name, target.get_base_layer()) + else: + # This is for backwards compatibility with PEFT <= 0.6.2. + # TODO can be removed once that PEFT version is no longer supported. + from peft.tuners.lora import LoraLayer + + for name, module in model.named_children(): + if len(list(module.children())) > 0: + ## compound module, go inside it + recurse_remove_peft_layers(module) + + module_replaced = False + + if isinstance(module, LoraLayer) and isinstance(module, torch.nn.Linear): + new_module = torch.nn.Linear( + module.in_features, + module.out_features, + bias=module.bias is not None, + ).to(module.weight.device) + new_module.weight = module.weight + if module.bias is not None: + new_module.bias = module.bias + + module_replaced = True + elif isinstance(module, LoraLayer) and isinstance(module, torch.nn.Conv2d): + new_module = torch.nn.Conv2d( + module.in_channels, + module.out_channels, + module.kernel_size, + module.stride, + module.padding, + module.dilation, + module.groups, + ).to(module.weight.device) + + new_module.weight = module.weight + if module.bias is not None: + new_module.bias = module.bias + + module_replaced = True + + if module_replaced: + setattr(model, name, new_module) + del module + + empty_device_cache() + return model + + +def scale_lora_layers(model, weight): + """ + Adjust the weightage given to the LoRA layers of the model. + + Args: + model (`torch.nn.Module`): + The model to scale. + weight (`float`): + The weight to be given to the LoRA layers. + """ + from peft.tuners.tuners_utils import BaseTunerLayer + + if weight == 1.0: + return + + for module in model.modules(): + if isinstance(module, BaseTunerLayer): + module.scale_layer(weight) + + +def unscale_lora_layers(model, weight: float | None = None): + """ + Removes the previously passed weight given to the LoRA layers of the model. + + Args: + model (`torch.nn.Module`): + The model to scale. + weight (`float`, *optional*): + The weight to be given to the LoRA layers. If no scale is passed the scale of the lora layer will be + re-initialized to the correct value. If 0.0 is passed, we will re-initialize the scale with the correct + value. + """ + from peft.tuners.tuners_utils import BaseTunerLayer + + if weight is None or weight == 1.0: + return + + for module in model.modules(): + if isinstance(module, BaseTunerLayer): + if weight != 0: + module.unscale_layer(weight) + else: + for adapter_name in module.active_adapters: + # if weight == 0 unscale should re-set the scale to the original value. + module.set_scale(adapter_name, 1.0) + + +# Maps the LoRA weight suffix of a text-encoder network alpha key (diffusers naming) +# to the corresponding PEFT module name suffix. PEFT's `alpha_pattern` keys must +# match the end of the PEFT module names (see `peft.utils.get_pattern_key`), e.g. +# the alpha key `text_model.encoder.layers.1.self_attn.to_k_lora.down.weight.alpha` +# belongs to the module `...self_attn.k_proj` -- not its parent block. +_TEXT_ENCODER_ALPHA_SUFFIX_TO_MODULE_SUFFIX = ( + (".to_q_lora.down.weight.alpha", ".q_proj"), + (".to_k_lora.down.weight.alpha", ".k_proj"), + (".to_v_lora.down.weight.alpha", ".v_proj"), + (".to_out_lora.down.weight.alpha", ".out_proj"), + (".lora_linear_layer.down.weight.alpha", ""), + # e.g. `text_model.text_projection.alpha` + (".alpha", ""), +) + + +def _text_encoder_alpha_pattern_key(alpha_key: str) -> str: + """Derive the PEFT `alpha_pattern` key from a text-encoder network alpha key.""" + for suffix, module_suffix in _TEXT_ENCODER_ALPHA_SUFFIX_TO_MODULE_SUFFIX: + if alpha_key.endswith(suffix): + key = alpha_key.removesuffix(suffix) + module_suffix + # The loader strips the `text_model.` prefix for flattened text encoders, + # so the pattern key must not require it: `get_pattern_key` matches at the + # end of the module name, making the shorter key match either way. + return key.removeprefix("text_model.") + return alpha_key + + +def get_peft_kwargs( + rank_dict, network_alpha_dict, peft_state_dict, is_unet=True, model_state_dict=None, adapter_name=None +): + rank_pattern = {} + alpha_pattern = {} + r = lora_alpha = list(rank_dict.values())[0] + has_alphas = network_alpha_dict is not None and len(network_alpha_dict) > 0 + + if len(set(rank_dict.values())) > 1: + # get the rank occurring the most number of times + r = collections.Counter(rank_dict.values()).most_common()[0][0] + + # for modules with rank different from the most occurring rank, add it to the `rank_pattern` + rank_pattern = dict(filter(lambda x: x[1] != r, rank_dict.items())) + rank_pattern = {k.split(".lora_B.")[0]: v for k, v in rank_pattern.items()} + + if not has_alphas: + # No alpha data in the checkpoint: the diffusers/PEFT convention is + # `W_eff = W + lora_B @ lora_A`, i.e. alpha == rank per module (scale 1.0). + # Mirror the ranks into the alphas so every module keeps scale 1.0. + lora_alpha = r + alpha_pattern = dict(rank_pattern) + + if network_alpha_dict is not None and len(network_alpha_dict) > 0: + if len(set(network_alpha_dict.values())) > 1: + # get the alpha occurring the most number of times + lora_alpha = collections.Counter(network_alpha_dict.values()).most_common()[0][0] + + # for modules with alpha different from the most occurring alpha, add it to the `alpha_pattern` + alpha_pattern = dict(filter(lambda x: x[1] != lora_alpha, network_alpha_dict.items())) + if is_unet: + alpha_pattern = { + ".".join(k.split(".lora_A.")[0].split(".")).replace(".alpha", ""): v + for k, v in alpha_pattern.items() + } + else: + alpha_pattern = {_text_encoder_alpha_pattern_key(k): v for k, v in alpha_pattern.items()} + else: + lora_alpha = set(network_alpha_dict.values()).pop() + + target_modules = list({name.split(".lora")[0] for name in peft_state_dict.keys()}) + use_dora = any("lora_magnitude_vector" in k for k in peft_state_dict) + # for now we know that the "bias" keys are only associated with `lora_B`. + lora_bias = any("lora_B" in k and k.endswith(".bias") for k in peft_state_dict) + + lora_config_kwargs = { + "r": r, + "lora_alpha": lora_alpha, + "rank_pattern": rank_pattern, + "alpha_pattern": alpha_pattern, + "target_modules": target_modules, + "use_dora": use_dora, + "lora_bias": lora_bias, + } + + return lora_config_kwargs + + +def get_adapter_name(model): + from peft.tuners.tuners_utils import BaseTunerLayer + + for module in model.modules(): + if isinstance(module, BaseTunerLayer): + return f"default_{len(module.r)}" + return "default_0" + + +def set_adapter_layers(model, enabled=True): + from peft.tuners.tuners_utils import BaseTunerLayer + + for module in model.modules(): + if isinstance(module, BaseTunerLayer): + # The recent version of PEFT needs to call `enable_adapters` instead + if hasattr(module, "enable_adapters"): + module.enable_adapters(enabled=enabled) + else: + module.disable_adapters = not enabled + + +def delete_adapter_layers(model, adapter_name): + from peft.tuners.tuners_utils import BaseTunerLayer + + for module in model.modules(): + if isinstance(module, BaseTunerLayer): + if hasattr(module, "delete_adapter"): + module.delete_adapter(adapter_name) + else: + raise ValueError( + "The version of PEFT you are using is not compatible, please use a version that is greater than 0.6.1" + ) + + # For transformers integration - we need to pop the adapter from the config + if getattr(model, "_hf_peft_config_loaded", False) and hasattr(model, "peft_config"): + model.peft_config.pop(adapter_name, None) + # In case all adapters are deleted, we need to delete the config + # and make sure to set the flag to False + if len(model.peft_config) == 0: + del model.peft_config + model._hf_peft_config_loaded = None + + +def set_weights_and_activate_adapters(model, adapter_names, weights): + from peft.tuners.tuners_utils import BaseTunerLayer + + def get_module_weight(weight_for_adapter, module_name): + if not isinstance(weight_for_adapter, dict): + # If weight_for_adapter is a single number, always return it. + return weight_for_adapter + + for layer_name, weight_ in weight_for_adapter.items(): + if layer_name in module_name: + return weight_ + + parts = module_name.split(".") + # e.g. key = "down_blocks.1.attentions.0" + key = f"{parts[0]}.{parts[1]}.attentions.{parts[3]}" + block_weight = weight_for_adapter.get(key, 1.0) + + return block_weight + + for module_name, module in model.named_modules(): + if isinstance(module, BaseTunerLayer): + # For backward compatibility with previous PEFT versions, set multiple active adapters + if hasattr(module, "set_adapter"): + module.set_adapter(adapter_names) + else: + module.active_adapter = adapter_names + + # Set the scaling weight for each adapter for this module + for adapter_name, weight in zip(adapter_names, weights): + module.set_scale(adapter_name, get_module_weight(weight, module_name)) + + +def apply_lora_scale(kwargs_name: str = "joint_attention_kwargs"): + """ + Decorator to automatically handle LoRA layer scaling/unscaling in forward methods. + + This decorator extracts the `lora_scale` from the specified kwargs parameter, applies scaling before the forward + pass, and ensures unscaling happens after, even if an exception occurs. + + Args: + kwargs_name (`str`, defaults to `"joint_attention_kwargs"`): + The name of the keyword argument that contains the LoRA scale. Common values include + "joint_attention_kwargs", "attention_kwargs", "cross_attention_kwargs", etc. + """ + + def decorator(forward_fn): + @functools.wraps(forward_fn) + def wrapper(self, *args, **kwargs): + from . import USE_PEFT_BACKEND + + lora_scale = 1.0 + attention_kwargs = kwargs.get(kwargs_name) + + if attention_kwargs is not None: + attention_kwargs = attention_kwargs.copy() + kwargs[kwargs_name] = attention_kwargs + lora_scale = attention_kwargs.pop("scale", 1.0) + + if not USE_PEFT_BACKEND and lora_scale != 1.0: + logger.warning( + f"Passing `scale` via `{kwargs_name}` when not using the PEFT backend is ineffective." + ) + + # Apply LoRA scaling if using PEFT backend + if USE_PEFT_BACKEND: + scale_lora_layers(self, lora_scale) + + try: + # Execute the forward pass + result = forward_fn(self, *args, **kwargs) + return result + finally: + # Always unscale, even if forward pass raises an exception + if USE_PEFT_BACKEND: + unscale_lora_layers(self, lora_scale) + + return wrapper + + return decorator + + +def require_peft_backend(fn): + """ + Decorator marking a method that needs the PEFT backend, i.e. compatible `peft` and `transformers` installations. + """ + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + from . import USE_PEFT_BACKEND + + if not USE_PEFT_BACKEND: + raise ValueError(f"PEFT backend is required for `{fn.__name__}()`.") + + return fn(*args, **kwargs) + + return wrapper + + +def check_peft_version(min_version: str) -> None: + r""" + Checks if the version of PEFT is compatible. + + Args: + version (`str`): + The version of PEFT to check against. + """ + if not is_peft_available(): + raise ValueError("PEFT is not installed. Please install it with `pip install peft`") + + is_peft_version_compatible = version.parse(importlib.metadata.version("peft")) >= version.parse(min_version) + + if not is_peft_version_compatible: + raise ValueError( + f"The version of PEFT you are using is not compatible, please use a version that is at least {min_version}" + ) + + +def _maybe_fuse_qkv_projections_for_lokr(model, state_dict) -> None: + """ + Fuse the model's QKV projections when a peft-format LoKr state dict targets fused `to_qkv` / `to_added_qkv` + projections that the model does not have yet. + + BFL-format Flux2 LoKr checkpoints apply LoKr to the fused QKV projections. Unlike a LoRA delta, a Kronecker product + delta over the fused projection cannot be split exactly into separate Q/K/V factors, so the model's projections are + fused instead and the adapter maps 1:1. + """ + fused_targets = { + module + for module in (k.rpartition(".lokr_")[0] for k in state_dict if ".lokr_" in k) + if module.rsplit(".", 1)[-1] in ("to_qkv", "to_added_qkv") + } + named_modules = dict(model.named_modules()) + if all(module in named_modules for module in fused_targets) or not hasattr(model, "fuse_qkv_projections"): + return + + if getattr(model, "is_quantized", False): + raise ValueError( + "This LoKr checkpoint targets fused QKV projections. Fusing concatenates the Q/K/V weights into a new " + "`nn.Linear`, which is not possible with quantized weights. Please load the transformer without " + "quantization." + ) + + # Fusing replaces to_q/to_k/to_v (and the add_*_proj) with a single projection, which would orphan any adapter + # already injected on the unfused ones. + from peft.tuners.tuners_utils import BaseTunerLayer + + unfused_projections = {"to_q", "to_k", "to_v", "add_q_proj", "add_k_proj", "add_v_proj"} + adapted = [ + name + for name, module in named_modules.items() + if isinstance(module, BaseTunerLayer) and name.rsplit(".", 1)[-1] in unfused_projections + ] + if adapted: + raise ValueError( + "This LoKr checkpoint targets fused QKV projections, but an adapter is already loaded on the unfused " + f"projections (e.g. `{adapted[0]}`). Unload it with `unload_lora_weights()` before loading this checkpoint." + ) + + logger.info("The LoKr checkpoint targets fused QKV projections; calling `fuse_qkv_projections()` on the model.") + model.fuse_qkv_projections() + + +def _create_lokr_config(state_dict, metadata): + """ + Create a `LoKrConfig` from a peft-format LoKr state dict (keys like `{module}.lokr_w1`). + + Without metadata, the config is inferred from the tensor shapes. The checkpoint alpha is expected to be already + baked into the weights by the state dict conversion, so `alpha` is set equal to the rank (runtime scaling 1.0). + peft re-derives each module's Kronecker factorization from `decompose_factor` and only creates rank-decomposed + factors when the rank is small compared to the factorized dimensions, so for modules whose checkpoint factors are + full matrices the rank is set to `max(lokr_w2.shape)` to make peft create full matrices as well. + """ + from peft import LoKrConfig + from peft.tuners.lokr.layer import factorization + + if metadata is not None: + try: + return LoKrConfig(**metadata) + except TypeError as e: + raise TypeError("`LoKrConfig` class could not be instantiated.") from e + + modules = sorted({k.rpartition(".lokr_")[0] for k in state_dict if ".lokr_" in k}) + + # Reconstruct each module's factorized dimensions, (out_l, out_k) x (in_m, in_n), from the checkpoint. A factor + # is either stored as a full matrix (`lokr_w1`) or rank-decomposed into `lokr_w1_a @ lokr_w1_b` (same for w2). + factorizations = {} + rank_dict = {} + for module in modules: + w1, w1_a, w1_b = (state_dict.get(f"{module}.lokr_w1{s}") for s in ("", "_a", "_b")) + w2, w2_a, w2_b = (state_dict.get(f"{module}.lokr_w2{s}") for s in ("", "_a", "_b")) + out_l, in_m = w1.shape if w1 is not None else (w1_a.shape[0], w1_b.shape[1]) + out_k, in_n = w2.shape if w2 is not None else (w2_a.shape[0], w2_b.shape[1]) + factorizations[module] = ((out_l, out_k), (in_m, in_n)) + if w2_a is not None: + rank_dict[module] = w2_a.shape[1] + elif w1_a is not None: + rank_dict[module] = w1_a.shape[1] + else: + rank_dict[module] = max(w2.shape) + + # Find the `decompose_factor` under which peft reproduces the checkpoint factorizations. A fixed factor shows up + # as the left dimension of the modules it divides (modules it does not divide fall back to a near-square + # factorization, like with factor -1), so every observed left dimension is a candidate. + left_dims = {dims[0][0] for dims in factorizations.values()} + decompose_factor = None + for candidate in sorted(left_dims) + [-1]: + if all( + factorization(out_l * out_k, candidate) == (out_l, out_k) + and factorization(in_m * in_n, candidate) == (in_m, in_n) + for (out_l, out_k), (in_m, in_n) in factorizations.values() + ): + decompose_factor = candidate + break + if decompose_factor is None: + raise ValueError( + "Could not infer a `decompose_factor` that reproduces the Kronecker factorizations of this LoKr " + "state dict. Please open an issue: https://github.com/huggingface/diffusers/issues/new" + ) + + r = collections.Counter(rank_dict.values()).most_common(1)[0][0] + rank_pattern = {k: v for k, v in rank_dict.items() if v != r} + + lokr_config_kwargs = { + "r": r, + "alpha": r, + "rank_pattern": rank_pattern, + "alpha_pattern": dict(rank_pattern), + "target_modules": modules, + "decompose_both": any(".lokr_w1_a" in k for k in state_dict), + "decompose_factor": decompose_factor, + } + try: + return LoKrConfig(**lokr_config_kwargs) + except TypeError as e: + raise TypeError("`LoKrConfig` class could not be instantiated.") from e + + +def _create_lora_config( + state_dict, network_alphas, metadata, rank_pattern_dict, is_unet=True, model_state_dict=None, adapter_name=None +): + from peft import LoraConfig + + if metadata is not None: + lora_config_kwargs = metadata + else: + lora_config_kwargs = get_peft_kwargs( + rank_pattern_dict, + network_alpha_dict=network_alphas, + peft_state_dict=state_dict, + is_unet=is_unet, + model_state_dict=model_state_dict, + adapter_name=adapter_name, + ) + + try: + return LoraConfig(**lora_config_kwargs) + except TypeError as e: + raise TypeError("`LoraConfig` class could not be instantiated.") from e + + +def _maybe_warn_for_unhandled_keys(incompatible_keys, adapter_name: str) -> None: + warn_msg = "" + if incompatible_keys is not None: + # Check only for unexpected keys. + unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None) + if unexpected_keys: + lora_unexpected_keys = [k for k in unexpected_keys if ".lora_" in k or "lokr_" in k] + if lora_unexpected_keys: + warn_msg = ( + f"Loading adapter weights from state_dict led to unexpected keys found in the model:" + f" {', '.join(lora_unexpected_keys)}. " + ) + + # Filter missing keys specific to the current adapter. + missing_keys = getattr(incompatible_keys, "missing_keys", None) + if missing_keys: + lora_missing_keys = [k for k in missing_keys if (".lora_" in k or "lokr_" in k) and adapter_name in k] + if lora_missing_keys: + warn_msg += ( + f"Loading adapter weights from state_dict led to missing keys in the model:" + f" {', '.join(lora_missing_keys)}." + ) + + if warn_msg: + logger.warning(warn_msg) From 1c614571c11efd017bd6ccdfee9916bbbf5a4b95 Mon Sep 17 00:00:00 2001 From: Akhil Gupta Date: Mon, 12 Oct 2026 21:12:13 -0700 Subject: [PATCH 4/5] Delete src/diffusers/utils/peft_utils.py --- src/diffusers/utils/peft_utils.py | 566 ------------------------------ 1 file changed, 566 deletions(-) delete mode 100644 src/diffusers/utils/peft_utils.py diff --git a/src/diffusers/utils/peft_utils.py b/src/diffusers/utils/peft_utils.py deleted file mode 100644 index 1096d5a2e3e8..000000000000 --- a/src/diffusers/utils/peft_utils.py +++ /dev/null @@ -1,566 +0,0 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -""" -PEFT utilities: Utilities related to peft library -""" - -import collections -import functools -import importlib - -from packaging import version - -from . import logging -from .import_utils import is_peft_available, is_torch_available -from .torch_utils import empty_device_cache - - -logger = logging.get_logger(__name__) - -if is_torch_available(): - import torch - - -def recurse_remove_peft_layers(model): - r""" - Recursively replace all instances of `LoraLayer` with corresponding new layers in `model`. - """ - from peft.tuners.tuners_utils import BaseTunerLayer - - has_base_layer_pattern = False - for module in model.modules(): - if isinstance(module, BaseTunerLayer): - has_base_layer_pattern = hasattr(module, "base_layer") - break - - if has_base_layer_pattern: - from peft.utils import _get_submodules - - key_list = [key for key, _ in model.named_modules() if "lora" not in key] - for key in key_list: - try: - parent, target, target_name = _get_submodules(model, key) - except AttributeError: - continue - if hasattr(target, "base_layer"): - setattr(parent, target_name, target.get_base_layer()) - else: - # This is for backwards compatibility with PEFT <= 0.6.2. - # TODO can be removed once that PEFT version is no longer supported. - from peft.tuners.lora import LoraLayer - - for name, module in model.named_children(): - if len(list(module.children())) > 0: - ## compound module, go inside it - recurse_remove_peft_layers(module) - - module_replaced = False - - if isinstance(module, LoraLayer) and isinstance(module, torch.nn.Linear): - new_module = torch.nn.Linear( - module.in_features, - module.out_features, - bias=module.bias is not None, - ).to(module.weight.device) - new_module.weight = module.weight - if module.bias is not None: - new_module.bias = module.bias - - module_replaced = True - elif isinstance(module, LoraLayer) and isinstance(module, torch.nn.Conv2d): - new_module = torch.nn.Conv2d( - module.in_channels, - module.out_channels, - module.kernel_size, - module.stride, - module.padding, - module.dilation, - module.groups, - ).to(module.weight.device) - - new_module.weight = module.weight - if module.bias is not None: - new_module.bias = module.bias - - module_replaced = True - - if module_replaced: - setattr(model, name, new_module) - del module - - empty_device_cache() - return model - - -def scale_lora_layers(model, weight): - """ - Adjust the weightage given to the LoRA layers of the model. - - Args: - model (`torch.nn.Module`): - The model to scale. - weight (`float`): - The weight to be given to the LoRA layers. - """ - from peft.tuners.tuners_utils import BaseTunerLayer - - if weight == 1.0: - return - - for module in model.modules(): - if isinstance(module, BaseTunerLayer): - module.scale_layer(weight) - - -def unscale_lora_layers(model, weight: float | None = None): - """ - Removes the previously passed weight given to the LoRA layers of the model. - - Args: - model (`torch.nn.Module`): - The model to scale. - weight (`float`, *optional*): - The weight to be given to the LoRA layers. If no scale is passed the scale of the lora layer will be - re-initialized to the correct value. If 0.0 is passed, we will re-initialize the scale with the correct - value. - """ - from peft.tuners.tuners_utils import BaseTunerLayer - - if weight is None or weight == 1.0: - return - - for module in model.modules(): - if isinstance(module, BaseTunerLayer): - if weight != 0: - module.unscale_layer(weight) - else: - for adapter_name in module.active_adapters: - # if weight == 0 unscale should re-set the scale to the original value. - module.set_scale(adapter_name, 1.0) - - -# Maps the LoRA weight suffix of a text-encoder network alpha key (diffusers naming) -# to the corresponding PEFT module name suffix. PEFT's `alpha_pattern` keys must -# match the end of the PEFT module names (see `peft.utils.get_pattern_key`), e.g. -# the alpha key `text_model.encoder.layers.1.self_attn.to_k_lora.down.weight.alpha` -# belongs to the module `...self_attn.k_proj` -- not its parent block. -_TEXT_ENCODER_ALPHA_SUFFIX_TO_MODULE_SUFFIX = ( - (".to_q_lora.down.weight.alpha", ".q_proj"), - (".to_k_lora.down.weight.alpha", ".k_proj"), - (".to_v_lora.down.weight.alpha", ".v_proj"), - (".to_out_lora.down.weight.alpha", ".out_proj"), - (".lora_linear_layer.down.weight.alpha", ""), - # e.g. `text_model.text_projection.alpha` - (".alpha", ""), -) - - -def _text_encoder_alpha_pattern_key(alpha_key: str) -> str: - """Derive the PEFT `alpha_pattern` key from a text-encoder network alpha key.""" - for suffix, module_suffix in _TEXT_ENCODER_ALPHA_SUFFIX_TO_MODULE_SUFFIX: - if alpha_key.endswith(suffix): - key = alpha_key[: -len(suffix)] + module_suffix - # The loader strips the `text_model.` prefix for flattened text encoders, - # so the pattern key must not require it: `get_pattern_key` matches at the - # end of the module name, making the shorter key match either way. - return key.removeprefix("text_model.") - return alpha_key - - -def get_peft_kwargs( - rank_dict, network_alpha_dict, peft_state_dict, is_unet=True, model_state_dict=None, adapter_name=None -): - rank_pattern = {} - alpha_pattern = {} - r = lora_alpha = list(rank_dict.values())[0] - has_alphas = network_alpha_dict is not None and len(network_alpha_dict) > 0 - - if len(set(rank_dict.values())) > 1: - # get the rank occurring the most number of times - r = collections.Counter(rank_dict.values()).most_common()[0][0] - - # for modules with rank different from the most occurring rank, add it to the `rank_pattern` - rank_pattern = dict(filter(lambda x: x[1] != r, rank_dict.items())) - rank_pattern = {k.split(".lora_B.")[0]: v for k, v in rank_pattern.items()} - - if not has_alphas: - # No alpha data in the checkpoint: the diffusers/PEFT convention is - # `W_eff = W + lora_B @ lora_A`, i.e. alpha == rank per module (scale 1.0). - # Mirror the ranks into the alphas so every module keeps scale 1.0. - lora_alpha = r - alpha_pattern = dict(rank_pattern) - - if network_alpha_dict is not None and len(network_alpha_dict) > 0: - if len(set(network_alpha_dict.values())) > 1: - # get the alpha occurring the most number of times - lora_alpha = collections.Counter(network_alpha_dict.values()).most_common()[0][0] - - # for modules with alpha different from the most occurring alpha, add it to the `alpha_pattern` - alpha_pattern = dict(filter(lambda x: x[1] != lora_alpha, network_alpha_dict.items())) - if is_unet: - alpha_pattern = { - ".".join(k.split(".lora_A.")[0].split(".")).replace(".alpha", ""): v - for k, v in alpha_pattern.items() - } - else: - alpha_pattern = {_text_encoder_alpha_pattern_key(k): v for k, v in alpha_pattern.items()} - else: - lora_alpha = set(network_alpha_dict.values()).pop() - - target_modules = list({name.split(".lora")[0] for name in peft_state_dict.keys()}) - use_dora = any("lora_magnitude_vector" in k for k in peft_state_dict) - # for now we know that the "bias" keys are only associated with `lora_B`. - lora_bias = any("lora_B" in k and k.endswith(".bias") for k in peft_state_dict) - - lora_config_kwargs = { - "r": r, - "lora_alpha": lora_alpha, - "rank_pattern": rank_pattern, - "alpha_pattern": alpha_pattern, - "target_modules": target_modules, - "use_dora": use_dora, - "lora_bias": lora_bias, - } - - return lora_config_kwargs - - -def get_adapter_name(model): - from peft.tuners.tuners_utils import BaseTunerLayer - - for module in model.modules(): - if isinstance(module, BaseTunerLayer): - return f"default_{len(module.r)}" - return "default_0" - - -def set_adapter_layers(model, enabled=True): - from peft.tuners.tuners_utils import BaseTunerLayer - - for module in model.modules(): - if isinstance(module, BaseTunerLayer): - # The recent version of PEFT needs to call `enable_adapters` instead - if hasattr(module, "enable_adapters"): - module.enable_adapters(enabled=enabled) - else: - module.disable_adapters = not enabled - - -def delete_adapter_layers(model, adapter_name): - from peft.tuners.tuners_utils import BaseTunerLayer - - for module in model.modules(): - if isinstance(module, BaseTunerLayer): - if hasattr(module, "delete_adapter"): - module.delete_adapter(adapter_name) - else: - raise ValueError( - "The version of PEFT you are using is not compatible, please use a version that is greater than 0.6.1" - ) - - # For transformers integration - we need to pop the adapter from the config - if getattr(model, "_hf_peft_config_loaded", False) and hasattr(model, "peft_config"): - model.peft_config.pop(adapter_name, None) - # In case all adapters are deleted, we need to delete the config - # and make sure to set the flag to False - if len(model.peft_config) == 0: - del model.peft_config - model._hf_peft_config_loaded = None - - -def set_weights_and_activate_adapters(model, adapter_names, weights): - from peft.tuners.tuners_utils import BaseTunerLayer - - def get_module_weight(weight_for_adapter, module_name): - if not isinstance(weight_for_adapter, dict): - # If weight_for_adapter is a single number, always return it. - return weight_for_adapter - - for layer_name, weight_ in weight_for_adapter.items(): - if layer_name in module_name: - return weight_ - - parts = module_name.split(".") - # e.g. key = "down_blocks.1.attentions.0" - key = f"{parts[0]}.{parts[1]}.attentions.{parts[3]}" - block_weight = weight_for_adapter.get(key, 1.0) - - return block_weight - - for module_name, module in model.named_modules(): - if isinstance(module, BaseTunerLayer): - # For backward compatibility with previous PEFT versions, set multiple active adapters - if hasattr(module, "set_adapter"): - module.set_adapter(adapter_names) - else: - module.active_adapter = adapter_names - - # Set the scaling weight for each adapter for this module - for adapter_name, weight in zip(adapter_names, weights): - module.set_scale(adapter_name, get_module_weight(weight, module_name)) - - -def apply_lora_scale(kwargs_name: str = "joint_attention_kwargs"): - """ - Decorator to automatically handle LoRA layer scaling/unscaling in forward methods. - - This decorator extracts the `lora_scale` from the specified kwargs parameter, applies scaling before the forward - pass, and ensures unscaling happens after, even if an exception occurs. - - Args: - kwargs_name (`str`, defaults to `"joint_attention_kwargs"`): - The name of the keyword argument that contains the LoRA scale. Common values include - "joint_attention_kwargs", "attention_kwargs", "cross_attention_kwargs", etc. - """ - - def decorator(forward_fn): - @functools.wraps(forward_fn) - def wrapper(self, *args, **kwargs): - from . import USE_PEFT_BACKEND - - lora_scale = 1.0 - attention_kwargs = kwargs.get(kwargs_name) - - if attention_kwargs is not None: - attention_kwargs = attention_kwargs.copy() - kwargs[kwargs_name] = attention_kwargs - lora_scale = attention_kwargs.pop("scale", 1.0) - - if not USE_PEFT_BACKEND and lora_scale != 1.0: - logger.warning( - f"Passing `scale` via `{kwargs_name}` when not using the PEFT backend is ineffective." - ) - - # Apply LoRA scaling if using PEFT backend - if USE_PEFT_BACKEND: - scale_lora_layers(self, lora_scale) - - try: - # Execute the forward pass - result = forward_fn(self, *args, **kwargs) - return result - finally: - # Always unscale, even if forward pass raises an exception - if USE_PEFT_BACKEND: - unscale_lora_layers(self, lora_scale) - - return wrapper - - return decorator - - -def require_peft_backend(fn): - """ - Decorator marking a method that needs the PEFT backend, i.e. compatible `peft` and `transformers` installations. - """ - - @functools.wraps(fn) - def wrapper(*args, **kwargs): - from . import USE_PEFT_BACKEND - - if not USE_PEFT_BACKEND: - raise ValueError(f"PEFT backend is required for `{fn.__name__}()`.") - - return fn(*args, **kwargs) - - return wrapper - - -def check_peft_version(min_version: str) -> None: - r""" - Checks if the version of PEFT is compatible. - - Args: - version (`str`): - The version of PEFT to check against. - """ - if not is_peft_available(): - raise ValueError("PEFT is not installed. Please install it with `pip install peft`") - - is_peft_version_compatible = version.parse(importlib.metadata.version("peft")) >= version.parse(min_version) - - if not is_peft_version_compatible: - raise ValueError( - f"The version of PEFT you are using is not compatible, please use a version that is at least {min_version}" - ) - - -def _maybe_fuse_qkv_projections_for_lokr(model, state_dict) -> None: - """ - Fuse the model's QKV projections when a peft-format LoKr state dict targets fused `to_qkv` / `to_added_qkv` - projections that the model does not have yet. - - BFL-format Flux2 LoKr checkpoints apply LoKr to the fused QKV projections. Unlike a LoRA delta, a Kronecker product - delta over the fused projection cannot be split exactly into separate Q/K/V factors, so the model's projections are - fused instead and the adapter maps 1:1. - """ - fused_targets = { - module - for module in (k.rpartition(".lokr_")[0] for k in state_dict if ".lokr_" in k) - if module.rsplit(".", 1)[-1] in ("to_qkv", "to_added_qkv") - } - named_modules = dict(model.named_modules()) - if all(module in named_modules for module in fused_targets) or not hasattr(model, "fuse_qkv_projections"): - return - - if getattr(model, "is_quantized", False): - raise ValueError( - "This LoKr checkpoint targets fused QKV projections. Fusing concatenates the Q/K/V weights into a new " - "`nn.Linear`, which is not possible with quantized weights. Please load the transformer without " - "quantization." - ) - - # Fusing replaces to_q/to_k/to_v (and the add_*_proj) with a single projection, which would orphan any adapter - # already injected on the unfused ones. - from peft.tuners.tuners_utils import BaseTunerLayer - - unfused_projections = {"to_q", "to_k", "to_v", "add_q_proj", "add_k_proj", "add_v_proj"} - adapted = [ - name - for name, module in named_modules.items() - if isinstance(module, BaseTunerLayer) and name.rsplit(".", 1)[-1] in unfused_projections - ] - if adapted: - raise ValueError( - "This LoKr checkpoint targets fused QKV projections, but an adapter is already loaded on the unfused " - f"projections (e.g. `{adapted[0]}`). Unload it with `unload_lora_weights()` before loading this checkpoint." - ) - - logger.info("The LoKr checkpoint targets fused QKV projections; calling `fuse_qkv_projections()` on the model.") - model.fuse_qkv_projections() - - -def _create_lokr_config(state_dict, metadata): - """ - Create a `LoKrConfig` from a peft-format LoKr state dict (keys like `{module}.lokr_w1`). - - Without metadata, the config is inferred from the tensor shapes. The checkpoint alpha is expected to be already - baked into the weights by the state dict conversion, so `alpha` is set equal to the rank (runtime scaling 1.0). - peft re-derives each module's Kronecker factorization from `decompose_factor` and only creates rank-decomposed - factors when the rank is small compared to the factorized dimensions, so for modules whose checkpoint factors are - full matrices the rank is set to `max(lokr_w2.shape)` to make peft create full matrices as well. - """ - from peft import LoKrConfig - from peft.tuners.lokr.layer import factorization - - if metadata is not None: - try: - return LoKrConfig(**metadata) - except TypeError as e: - raise TypeError("`LoKrConfig` class could not be instantiated.") from e - - modules = sorted({k.rpartition(".lokr_")[0] for k in state_dict if ".lokr_" in k}) - - # Reconstruct each module's factorized dimensions, (out_l, out_k) x (in_m, in_n), from the checkpoint. A factor - # is either stored as a full matrix (`lokr_w1`) or rank-decomposed into `lokr_w1_a @ lokr_w1_b` (same for w2). - factorizations = {} - rank_dict = {} - for module in modules: - w1, w1_a, w1_b = (state_dict.get(f"{module}.lokr_w1{s}") for s in ("", "_a", "_b")) - w2, w2_a, w2_b = (state_dict.get(f"{module}.lokr_w2{s}") for s in ("", "_a", "_b")) - out_l, in_m = w1.shape if w1 is not None else (w1_a.shape[0], w1_b.shape[1]) - out_k, in_n = w2.shape if w2 is not None else (w2_a.shape[0], w2_b.shape[1]) - factorizations[module] = ((out_l, out_k), (in_m, in_n)) - if w2_a is not None: - rank_dict[module] = w2_a.shape[1] - elif w1_a is not None: - rank_dict[module] = w1_a.shape[1] - else: - rank_dict[module] = max(w2.shape) - - # Find the `decompose_factor` under which peft reproduces the checkpoint factorizations. A fixed factor shows up - # as the left dimension of the modules it divides (modules it does not divide fall back to a near-square - # factorization, like with factor -1), so every observed left dimension is a candidate. - left_dims = {dims[0][0] for dims in factorizations.values()} - decompose_factor = None - for candidate in sorted(left_dims) + [-1]: - if all( - factorization(out_l * out_k, candidate) == (out_l, out_k) - and factorization(in_m * in_n, candidate) == (in_m, in_n) - for (out_l, out_k), (in_m, in_n) in factorizations.values() - ): - decompose_factor = candidate - break - if decompose_factor is None: - raise ValueError( - "Could not infer a `decompose_factor` that reproduces the Kronecker factorizations of this LoKr " - "state dict. Please open an issue: https://github.com/huggingface/diffusers/issues/new" - ) - - r = collections.Counter(rank_dict.values()).most_common(1)[0][0] - rank_pattern = {k: v for k, v in rank_dict.items() if v != r} - - lokr_config_kwargs = { - "r": r, - "alpha": r, - "rank_pattern": rank_pattern, - "alpha_pattern": dict(rank_pattern), - "target_modules": modules, - "decompose_both": any(".lokr_w1_a" in k for k in state_dict), - "decompose_factor": decompose_factor, - } - try: - return LoKrConfig(**lokr_config_kwargs) - except TypeError as e: - raise TypeError("`LoKrConfig` class could not be instantiated.") from e - - -def _create_lora_config( - state_dict, network_alphas, metadata, rank_pattern_dict, is_unet=True, model_state_dict=None, adapter_name=None -): - from peft import LoraConfig - - if metadata is not None: - lora_config_kwargs = metadata - else: - lora_config_kwargs = get_peft_kwargs( - rank_pattern_dict, - network_alpha_dict=network_alphas, - peft_state_dict=state_dict, - is_unet=is_unet, - model_state_dict=model_state_dict, - adapter_name=adapter_name, - ) - - try: - return LoraConfig(**lora_config_kwargs) - except TypeError as e: - raise TypeError("`LoraConfig` class could not be instantiated.") from e - - -def _maybe_warn_for_unhandled_keys(incompatible_keys, adapter_name: str) -> None: - warn_msg = "" - if incompatible_keys is not None: - # Check only for unexpected keys. - unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None) - if unexpected_keys: - lora_unexpected_keys = [k for k in unexpected_keys if ".lora_" in k or "lokr_" in k] - if lora_unexpected_keys: - warn_msg = ( - f"Loading adapter weights from state_dict led to unexpected keys found in the model:" - f" {', '.join(lora_unexpected_keys)}. " - ) - - # Filter missing keys specific to the current adapter. - missing_keys = getattr(incompatible_keys, "missing_keys", None) - if missing_keys: - lora_missing_keys = [k for k in missing_keys if (".lora_" in k or "lokr_" in k) and adapter_name in k] - if lora_missing_keys: - warn_msg += ( - f"Loading adapter weights from state_dict led to missing keys in the model:" - f" {', '.join(lora_missing_keys)}." - ) - - if warn_msg: - logger.warning(warn_msg) From 42d5e678d229a4e817b8f3bed4b833528325e38c Mon Sep 17 00:00:00 2001 From: Akhil Gupta Date: Mon, 12 Oct 2026 21:12:51 -0700 Subject: [PATCH 5/5] use removesuffix per review feedback --- .../utils/{diffusers14972-peft-utils.py => peft_utils.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename src/diffusers/utils/{diffusers14972-peft-utils.py => peft_utils.py} (100%) diff --git a/src/diffusers/utils/diffusers14972-peft-utils.py b/src/diffusers/utils/peft_utils.py similarity index 100% rename from src/diffusers/utils/diffusers14972-peft-utils.py rename to src/diffusers/utils/peft_utils.py