Fuse loras

[patrickvonplaten]

What does this PR do?

Great idea from @williamberman to allow fusing lora weights into the original weights

You can try it out with:

from diffusers import DiffusionPipeline
import torch

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.fuse_lora()

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]

You can reverse the effect of fuse_lora() by calling pipe.unfuse_lora(). Refer to the test cases to have a better handle on the implications.

TODOS:

• Conv layers
• Attention linear layers
• Text encoder LoRA fusion

BTW @sayakpaul we should probably refactor the attention lora processors to completely remove them and instead just work with the compatible linear layer class in the other attention processors (we'll have to do this anyways for the peft refactor) Done in #4765 by @patrickvonplaten.

[sayakpaul]

BTW @sayakpaul we should probably refactor the attention lora processors to completely remove them and instead just work with the compatible linear layer class in the other attention processors (we'll have to do this anyways for the peft refactor)

Could you please elaborate? I didn't get it at all.

[HuggingFaceDocBuilderDev]

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

[apolinario]

Kohya has a similar functionality and they have this ratio param: https://github.com/kohya-ss/sd-scripts/blob/main/networks/merge_lora.py#L33

python networks\merge_lora.py --sd_model ..\model\model.ckpt --save_to ..\lora_train1\model-char1-merged.safetensors --models ..\lora_train1\last.safetensors ..\lora_train2\last.safetensors --ratios 0.8 0.5

Would it make sense we enable this weighting into the (single) loading too we are doing too? 

[apolinario]

Also as mentioned internally on Slack, would be great to have a pipe.unet.unfuse_lora() for use-cases where on infra one may want to keep a model warm and swap out LoRAs

[sayakpaul]

@apolinario FWIW, this PR doesn't concern loading multiple LoRAs and fusing them. It simply fuses the LoRA params and the original model params in a way that doesn't affect the rest of our attention processing mechanism.

We already have issue requests for multiple LoRAs on GH.

[apolinario]

@sayakpaul sorry, I was not suggesting that we should tackle loading multiple LoRAs in this PR, I just used the loading multiple LoRAs code example of kohya to showcase the ratios param which was the thing I wanted to ask if we should do smth analogous or not on this PR

I have edited my comment for clarity

[sayakpaul]

IIUC, our scale parameter can be used to control how much the LoRA parama are being merged to the corresponding original params. What am I missing out on?

[apolinario]

Yup you are right, sounds redundant then if there's a general/global scale param for loading the LoRA

Edit: as discussed offline, while we do have a cross_attention_kwargs={"scale": 0.5} as stated here, that happens during inference, so in this PR we would need some sort of a scale param for the model merging

[patrickvonplaten]

BTW @sayakpaul we should probably refactor the attention lora processors to completely remove them and instead just work with the compatible linear layer class in the other attention processors (we'll have to do this anyways for the peft refactor)

Could you please elaborate? I didn't get it at all.

For this PR we need to add a specific fuse_lora function to

diffusers/src/diffusers/models/attention_processor.py

Line 1140 in aef11cb

class LoRAXFormersAttnProcessor(nn.Module):

and

diffusers/src/diffusers/models/attention_processor.py

Line 1254 in aef11cb

class LoRAAttnProcessor2_0(nn.Module):

to get it working. That's ok for now. However in the midterm we should try to fully delete these two classes and instead adapt:

diffusers/src/diffusers/models/attention_processor.py

Line 142 in aef11cb

self.to_k = nn.Linear(cross_attention_dim, inner_dim, bias=bias)

to

diffusers/src/diffusers/models/lora.py

Line 107 in aef11cb

class LoRACompatibleLinear(nn.Linear):

(this will need some changes in the loading and training of LoRA then, but we can make it work with fully backwards compatibility I believe - also cc @williamberman )

[patrickvonplaten]

@sayakpaul do you want to give this PR a try or do you prefer me to finish it?

[sayakpaul]

Giving it a try.

[sayakpaul]

@patrickvonplaten up for a review here.

[sayakpaul]

With this benchmarking script, I get the following on a V100:

{'fuse': False, 'total_time (ms)': '95874.1', 'memory (mb)': 13572}

{'fuse': True, 'total_time (ms)': '83744.8', 'memory (mb)': 13543}

[apolinario]

Two questions:

• Are we supporting a scale param for fusing the LoRA as previously discussed? Or would cross_attention_kwargs={"scale": 0.5} still work with fused LoRAs? 
• What happens if I fuse two LoRAs and then call unfuse_lora? Does it go back to the original state?

[sayakpaul]

I am observing some inconsistencies in the tests and unfuse_lora(). I will update the PR once it's ready for review.

[patrickvonplaten]

Ah I've found the problem I think. The problem is that we currently don't fuse the text encoder weights, when calling:

unload_lora_weights()

this line:

diffusers/src/diffusers/loaders.py

Line 1727 in 53f2e74

self._remove_text_encoder_monkey_patch()

is hit and causes the inference to be different.

The PR as it's currently already works correctly for non-textencoder LoRAs such as: https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_offset_example-lora_1.0.safetensors

you can try by doing:

from diffusers import DiffusionPipeline
import torch

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.unet.fuse_lora()
#This should have no effect as the LoRA is fused to the `unet` already, but instead it is removing the LoRA's effect
pipe.unload_lora_weights()

pipe.to(torch_dtype=torch.float16)
pipe.to("cuda")

torch.manual_seed(0)
prompt = "a mecha robot"
image = pipe(prompt, num_inference_steps=25).images[0]
image

I think we can actually also also fuse the text encoder weights no? Don't think that should be too difficult - what do you think @sayakpaul ?

[patrickvonplaten]

Apart from this the PR looks very nice! Tests that have been added are great

[sayakpaul]

@apolinario https://gist.github.com/sayakpaul/cd0395669002ae82634e57e5d26cc0e0.

Regular:

With fused LoRA:

With unload_lora_weights() (no-op):

With unfuse_lora():

[sayakpaul]

@patrickvonplaten PR is ready for a review. Would be great if you could do some testing as well.

My only question at this point is how do we allow for scale when the fuse_lora() has been called. We need to think about that and once we have a way, we should write thorough test cases.

[patrickvonplaten]

Very cool! Everything works fine and the implementation is clean - nice job!

Yes I think you touch upon a good point here regarding "lora_scale" - we should probably handle this in a new PR.
We currently have a also this 🚧 regarding the "lora_scale": #4751

=> I'd suggest to do the following in a follow-up PR:

• a) Make sure that "lora_scale" is applied to all LoRA layers. Instead of passing a cross_attention_kwargs through the attention processors, it's probably better to work with a "setter"-method here that iterates through all layers and if the layer is a LoRA layer it just sets the "lora_scale" for the specific class
• b) Also solve:

My only question at this point is how do we allow for scale when the fuse_lora() has been called. We need to think about that and once we have a way, we should write thorough test cases.
=> I think we can do this by just multiplying the fused weight with that "lora_scale" and it should work

• c) We should also through a warning if a user fuses more than one LoRA IMO

Feel free to take this PR - otherwise happy to take a look myself :-)

Merging this one now!

[apolinario]

Very nice @sayakpaul, however the robot after the pipe.unfuse_lora() is different than the robot before fusing the LoRA:

From your own gist and example (and I reproduced it here):

Before fusing the LoRA	After fusing and unfusing the LoRA
prompt: a mecha robot - seed: 0	prompt: a mecha robot - seed: 0

Is this expected? Do we know why? I think this can affect hot-unloading use-cases as after pipe.unfuse_lora() my understanding is that users should expect the same image as if the LoRA had never been fused

[patrickvonplaten]

Actually the reason unfuse gives different results might be because we don't do the computation in full fp32 precision here. In the _fuse_lora function we do the computation in full fp32 precision, but here we don't I think. Can we try to make sure that:

unfused_weight = fused_weight - torch.bmm(self.w_up[None, :], self.w_down[None, :])[0]

is always computed in full fp32 precision and only then we lower it potentially to fp16 dtype?

Guess we could also easily check this by casting the whole model to fp32 before doing fuse and unfuse to check cc @apolinario

[apolinario]

This is a good hypothesis but would this explain the behaviour of this happening in some LoRAs but not others, and keeping some residual style?

1. davizca87/vulcan → different image after unloading
2. ostris/crayon_style_lora_sdxl → same image after unloading
3. davizca87/sun-flower → different image after unloading (and also different than the image unlaoded in 1)
4. TheLastBen/Papercut_SDXL→ same image after unloading

[apolinario]

Failed cases ❌

davizca87/sun-flower seems to keep some of the sunflower vibe in the background

original generation	with lora fused	after unfusing (residual effects from the LoRA?)

davizca87/vulcan seems to keep some of the vulcan style in the outlines of the robot

original generation	with lora fused	after unfusing (residual effects from the LoRA?)

Success cases ✅

ostris/crayon_style_lora_sdxl seems to produce a perceptually identical image after unfusing

original generation	with lora fused	after unfusing (same as original generation)

TheLastBen/Papercut_SDXL and nerijs/pixel-art-xl also exhibit the same correct behavior

[sayakpaul]

I personally always worry about how numerical precision stems through a network and affect the end results. Have seen enough cases because of this to not sleep well at night. So, I will start with what @patrickvonplaten suggested.

FWIW, though, there's actually a fast test that ensures the following trip doesn't have side-effects:

load_lora_weights() -> fuse_lora() -> unload_lora_weights() gives you the outputs you would expect after doing fuse_lora().

Let me know if anything is unclear.

[apolinario]

Very nice! I guess an analogous test could be made to address what I reported for a future PR

Namely asserting that the two generate with no LoRA are matching:
generate with no LoRA before unfusing → load_lora_weights() → fuse_lora() → unfuse_lora() → generate with no LoRA after unfusing

This is the workflow I've reported above, the unfused unet seemingly still contains somewhat of a residue of the lora

[patrickvonplaten]

Should be fixed in: #4833 . Was a tricky issue that was caused by the patched text encoder LoRA layers being fully removed when doing unload_lora and therefore loosing their ability to unfuse. Hope it was ok to help out a bit here @sayakpaul

[sayakpaul]

Had started fuse-lora-pt2 and mentioned in Slack that I am looking into it. But okay.

[patrickvonplaten]

Very cool! Everything works fine and the implementation is clean - nice job!

Yes I think you touch upon a good point here regarding "lora_scale" - we should probably handle this in a new PR. We currently have a also this 🚧 regarding the "lora_scale": #4751

=> I'd suggest to do the following in a follow-up PR:

• a) Make sure that "lora_scale" is applied to all LoRA layers. Instead of passing a cross_attention_kwargs through the attention processors, it's probably better to work with a "setter"-method here that iterates through all layers and if the layer is a LoRA layer it just sets the "lora_scale" for the specific class
• b) Also solve:

My only question at this point is how do we allow for scale when the fuse_lora() has been called. We need to think about that and once we have a way, we should write thorough test cases.
=> I think we can do this by just multiplying the fused weight with that "lora_scale" and it should work

• c) We should also through a warning if a user fuses more than one LoRA IMO

Feel free to take this PR - otherwise happy to take a look myself :-)

Merging this one now!

@sayakpaul if you have some time to look into this before next week that would be incredible! This way we can make a super nice SDXL LoRA release :-)

[sayakpaul]

@sayakpaul if you have some time to look into this before next week that would be incredible! This way we can make a super nice SDXL LoRA release :-)

Prioritizing it.

[xhinker]

I see. My understanding was that we would get that for "free" - my understanding was, once a LoRA is fused to the Unet, loading another LoRA would be the same as loading it LoRA to a base model (given it's loading to a fused Unet). What am I missing?

@apolinario your understanding is correct. However, I think we need to think about that design a bit to not break the API consistency and consider the repercussions of that. I am not suggesting we shouldn't allow it, of course, we should. But the API and/or docs for it need to be very tight to not introduce any confusion for the users.

I agree! We should be careful about introducing new APIs here. A suggestion here would be to do the following: There is no harm in allowing the user to fuse multiple LoRAs, but as soon as two LoRAs are fused, the unfuse LoRA function doesn't work anymore. => What do you think about the idea of keeping an internal counter of how many LoRAs are fused and when this number is > 1, we throw an error when doing unfuse LoRA that states that multiple LoRAs have been fused and that the user needs to reload the whole model instead

I'd say only we have a full PEFT integration, we should have full support for fusing and unfusing multiple LoRAs by name.

Thoughts?

@patrickvonplaten @sayakpaul how about adding a LoRA state dictionary to remember the LoRAs and Scale Value. for example:

lora_state = [
     {
           'lora_name':"LoRA1",
           'lora_scale':0.7
     },
     {
           'lora_name':"LoRA2",
           'lora_scale':0.5
     }
]

When a new LoRA is added into the pipeline, update the scale delta to pipeline if the LoRA exists already, and update the lora_state object at the same time, add a new record if never added before,

Now, whenever user call the unfuse() function, the pipeline will remove the LoRA weights times the lora_state's lora_scale. In theory, it should work, thoughts?