[core / PEFT / LoRA] Integrate PEFT into Unet

[younesbelkada]

What does this PR do?

Second (actually third) step of PEFT integration into diffusers, this time PEFT is integrated on the Unet

The testing script is still the same, works with and without scale as expected

from diffusers import StableDiffusionPipeline
import os
import torch

path = "runwayml/stable-diffusion-v1-5"
lora_id = "takuma104/lora-test-text-encoder-lora-target"

pipe = StableDiffusionPipeline.from_pretrained(path, torch_dtype=torch.float16)
pipe.load_lora_weights(lora_id)
pipe = pipe.to("cuda")

prompt = "a red sks dog"

images = pipe(
    prompt=prompt, 
    num_inference_steps=15, 
    cross_attention_kwargs={"scale": 0.5},
    generator=torch.manual_seed(0)
).images


for i, image in enumerate(images):
    file_name = f"aa_{i}"
    os.makedirs("images-integration", exist_ok=True)
    path = os.path.join("images-integration", f"{file_name}-new-peft-2.png")
    image.save(path)

and this should return:

TODOs:

• deal with failing tests
• add tests
• think of a better way to deal with scaling / unscaling (decorator, ..?) inside call
• add fuse-unfuse tests for unet + TE
• add multi-adapter tests for text encoder
• add multi-adapter tests for text encoder + unet
• test it with Kohya checkpoints
• add deprecation messages inside LoRACompatiblexxx ?

cc @sayakpaul @patrickvonplaten @pacman100 @BenjaminBossan

[HuggingFaceDocBuilderDev]

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

[patrickvonplaten]

@younesbelkada let me know if you'd like me to do a review here

[younesbelkada]

Thanks @patrickvonplaten indeed I have a question here, since the changes made in pipeline_xxx needs to be copied over all the pipeline files I was wondering which approach looks best. @BenjaminBossan suggested offline that the second approach is better - which I also agree - let me know what do you think!

[sayakpaul]

I slightly prefer the second approach as I am not a big fan of using context managers for these kinds of use cases. I think it's better to be a bit verbose here.

[younesbelkada]

Thanks for all the reviews!
Regarding the test I have modified, if huggingface/peft#1017 gets merged, the test will pass on CPU (the test are used to ran on CPU). Since that test used to pass on CPU I think that usecase should be supported to avoid potential regressions.

The reason it used to pass is simply because with the old backend lora weights are always upcasted in fp32: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/lora.py#L257
Otherwise torch.bmm would fail:

>>> import torch
>>> t1 = torch.randn(1, 1).to(torch.float16)
>>> t2 = torch.randn(1, 1).to(torch.float16)
>>> torch.bmm(t1[None, :], t2[None, :])
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: "bmm" not implemented for 'Half'

[BenjaminBossan]

Thanks, this looks really good from my POV. I still left a few comments but those can be safely ignored. I also still think that get_list_adapters could be a confusing name to users, maybe you can find something better.

[BenjaminBossan]

This would have the same effect in two lines, but your version is probably easier to understand, just wanted to throw this out there.

Suggested change

	if not is_model_cpu_offload:
	is_model_cpu_offload = isinstance(component._hf_hook, CpuOffload)
	if not is_sequential_cpu_offload:
	is_sequential_cpu_offload = isinstance(component._hf_hook, AlignDevicesHook)
	is_model_cpu_offload |= isinstance(component._hf_hook, CpuOffload)
	is_sequential_cpu_offload |= isinstance(component._hf_hook, AlignDevicesHook)

[sayakpaul]

Would prefer the current implementation as is in favor of easier readability.

[sayakpaul]

Looks good to me, barring the pending comments. Thanks so much for iterating!

[patrickvonplaten]

Let's merge the PR and make sure we monitor the fast & slow test here

[patrickvonplaten]

Great job @younesbelkada and everybody involved here!

[younesbelkada]

Nice, thanks!
The slow CUDA lora old backend tests seem to pass: https://github.com/huggingface/diffusers/actions/runs/6509717155/job/17681829144 however the slow lora new backend are not run I think. I am happy to add them in the workflow, do you want me to open a PR for it?

[patrickvonplaten]

Nice, thanks! The slow CUDA lora old backend tests seem to pass: https://github.com/huggingface/diffusers/actions/runs/6509717155/job/17681829144 however the slow lora new backend are not run I think. I am happy to add them in the workflow, do you want me to open a PR for it?

Yes this would be great cc @DN6 here