feat: implement apply_lora_scale to remove boilerplate.

[sayakpaul]

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]

@DN6 would love to know your thoughts!

[sayakpaul]

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.

[DN6]

Nice! 👍🏽 Design looks good to me.

[sayakpaul]

@DN6 a gentle ping

[DN6]

LGTM 👍🏽 Left a comment on a small check.

[DN6]

This path won't ever be accessed since to reach here attention_kwargs has to be None.

I think it's supposed to be

  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."
          )

[sayakpaul]

Merging as discussed with Dhruv.