Skip to content

Conversation

@sayakpaul
Copy link
Member

What does this PR do?

Currently, we have this pattern throughout the modeling implementations:

if joint_attention_kwargs is not None:
    joint_attention_kwargs = joint_attention_kwargs.copy()
    lora_scale = joint_attention_kwargs.pop("scale", 1.0)
else:
    lora_scale = 1.0

if USE_PEFT_BACKEND:
    # weight the lora layers by setting `lora_scale` for each PEFT layer
    scale_lora_layers(self, lora_scale)
else:
    if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
        logger.warning(
            "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
        )


...


if USE_PEFT_BACKEND:
    # remove `lora_scale` from each PEFT layer
    unscale_lora_layers(self, lora_scale)

IMO, this is not pretty and should possibly be minimized for a clean and self-contained implementation of the forward().

Hence, this PR introduces a decorator apply_lora_scale that can be used to decorate the forward method of a model supporting LoRA. I think this will help us reduce a bunch of boilerplate code.

For keeping the PR simple, I have only applied the decorator to src/diffusers/models/transformers/transformer_flux.py The LoRA tests for it are passing, indicating this direction might be a nice one.

@sayakpaul sayakpaul requested a review from DN6 January 19, 2026 04:55
@sayakpaul
Copy link
Member Author

@DN6 would love to know your thoughts!


from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import FluxTransformer2DLoadersMixin, FromOriginalModelMixin, PeftAdapterMixin
from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
Copy link
Member Author

Choose a reason for hiding this comment

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

This yields 21 LoC deletions. We have this pattern in about 32 files. So, this amounts for a 672 deletions. Not bad, IMO.

@HuggingFaceDocBuilderDev

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants