[LoRA Attn Processors] Refactor LoRA Attn Processors

[patrickvonplaten]

What does this PR do?

This PR deprecates all the "LoRA..." attention processors as explained here: #4473 (comment)

This is due to the following reasons:

• Now that we have the concept of LoRACompatibleLinear there is no reason to not also leverage it in the attention classes. This way we have a single point of logic, removing overhead.
• It significantly helps when applying global LoRA functionality such as fuse_lora as seen here:
  Fuse loras #4473
• We reduce the number of attention processor classes significantly making the file easier to read

TODO:

• Refactor attn_processor method to (for now) return derprecated attn processor if LoRA layers are activated
• Make sure the saving format stays the same
• Make sure all LoRA training examples work
• Fix failing test

The PR works now for inference:

pipe=DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16)
pipe.load_lora_weights("stabilityai/stable-diffusion-xl-base-1.0", weight_name="sd_xl_offset_example-lora_1.0.safetensors")
pipe.to(torch_dtype=torch.float16)
pipe.to("cuda")
torch.manual_seed(0)
prompt="beautiful scenery nature glass bottle landscape, , purple galaxy bottle"negative_prompt="text, watermark"image=pipe(prompt, negative_prompt=negative_prompt, num_inference_steps=25).images[0]

@sayakpaul@williamberman I'd be very happy if you could give this a first review and if ok for you I can finish the final TODOs tomorrow.

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[patrickvonplaten]

Necessary for #4473

[sayakpaul]

Did some first pass and left comments. At this point, the changes seem a little too complicated to me but will probably have a better idea once the PR is in a more comprehensive state.

[patrickvonplaten]

A bit more explanation for what's going on here:

• a) Eventually we want to fully get rid of "lora" logic in diffusers and offload all this responsability to PEFT. The only thing that will be left in diffusers then will be {AnyLoraFormat}->PEFT_Format conversion logic which probably only amounts to renaming keys and other inputs.
• b) In this PR in a first step we want to unify the LoRA logic all into LoRACompatibleLinear and essentially deprecate all Attn Processor LoRAs. After this PR, all the computation logic will flow through LoRACompatibleLinear and we keep the LoRA attention processors only as a serialization format to not break backwards compatibility. Hence we need to ensure that calling unet.attn_processors will still return LoRA attention processors so that all the serialization format stays intact.
• c) Running the forward pass of any LoRA attention processor is already now deprecated. The LoRA attention processor serialization format will bit by bit be deprecated, but this should be done as part of the PEFT integration.

Overall this PR should make the whole LoRA loading and functioning easier to understand and build upon (such as needed for #4473

[patrickvonplaten]

1e-3 is a bit too tight for integration tests. We should be able to run most tests on many different GPUs without having them fail.

[sayakpaul]

You meant 1e-4.

[patrickvonplaten]

Updating the expected values here because I'm quite sure that this PR corrects a bug.
The lora:

 lora_model_id = "hf-internal-testing/sdxl-0.9-kamepan-lora"
lora_filename = "kame_sdxl_v2-000020-16rank.safetensors"

is quite unique since it has different network_alphas for q,k,v and out LoRA of the attention layer. Previously we always set the same network_alpha for all q,k,v and out (see here)

However, it might very well be that q,k,v and out have different network_alpha values. This is now possible as of this PR and therefore should mean it's a bug fix since this checkpoint has different network_alphas for each q,k,v and out.

[sayakpaul]

diffusers/src/diffusers/loaders.py

Line 1407 in 3bba44d

forname, attn_moduleintext_encoder_attn_modules(text_encoder):

suggests we assign different alphas to the modules you mentioned. What am I missing out on?

[sayakpaul]

Ah got it. We do that for the text encoder only it seems but not the UNet, is my understanding correct?

[patrickvonplaten]

Yeah exactly! I'm pretty sure it's corrected now :-)

[williamberman]

nice nice

[patrickvonplaten]

As soon as you're good with the PR @sayakpaul - I think we can merge this one and unblock the fuse_lora PR :-)

[sayakpaul]

Maybe we should just display list(is_lora_activated.keys()) to the end user.

[sayakpaul]

Neat.

I am a little concerned about:

and self.to_q.lora_layer is not None

i.e., only check if to_q.lora_layer is not None. Is there a better alternative?

[sayakpaul]

This bit is nasty but needed for ensuring backward compatibility.

[sayakpaul]

Do you want to directly throw a deprecation message when the class is initialized or would that be too brutal?

[patrickvonplaten]

I thought about this too - think this will be a bit too much though at the moment because it would mean every time we call unet.attn_processors we would get a deprecation warning which would be every time we save LoRAs. I'd suggest to only start throwing aggressive deprecation warnings when we do the peft integration

[sayakpaul]

Likewise.

[sayakpaul]

Like wise.

[sayakpaul]

Hard work. Thanks a lot.

Left some minor design-related comments but I don't think they block this PR from getting merged. So, I am gonna merge this to be able to focus on #4473.