[LoRA] allow fine-tuning of the text encoder with LoRA (using `peft`)

[sayakpaul]

We have had many requests (rightfully so) for allowing to fine-tune the text encoder with LoRA (such as #2683). This is quite useful for improving the quality of the generated samples. This issue thread aims to discuss a solution candidate that combines our LoRA attention processors and peft.

Proposed solution candidate

We stick to using AttnProcsLayers and LoRAAttnProcessor for handling the UNet:

from diffusers.loaders import AttnProcsLayers
from diffusers.models.attention_processor import LoRAAttnProcessor
from diffusers import UNet2DConditionModel

model_id = "runwayml/stable-diffusion-v1-5"
output_dir = "demo-peft"

unet = UNet2DConditionModel.from_pretrained(
    model_id, subfolder="unet", revision=None
)

lora_attn_procs = {}
for name in unet.attn_processors.keys():
    cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
    if name.startswith("mid_block"):
        hidden_size = unet.config.block_out_channels[-1]
    elif name.startswith("up_blocks"):
        block_id = int(name[len("up_blocks.")])
        hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
    elif name.startswith("down_blocks"):
        block_id = int(name[len("down_blocks.")])
        hidden_size = unet.config.block_out_channels[block_id]

    lora_attn_procs[name] = LoRAAttnProcessor(hidden_size=hidden_size, cross_attention_dim=cross_attention_dim)

unet.set_attn_processor(lora_attn_procs)
unet.save_attn_procs(output_dir)

For the text encoder, we do:

from peft import LoraConfig, get_peft_model
from transformers import CLIPTextModel

text_encoder = CLIPTextModel.from_pretrained(
    model_id, subfolder="text_encoder", revision=None
)

config = LoraConfig(
    r=4,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.0,
    bias="none",
)
lora_text_encoder = get_peft_model(text_encoder, config)
lora_text_encoder.save_pretrained(output_dir)

Now, if we do:

ls -lh {output_dir}

It should give us:

total 3.8M
-rw-r--r-- 1 root root  363 Mar 17 05:28 adapter_config.json
-rw-r--r-- 1 root root 593K Mar 17 05:28 adapter_model.bin
-rw-r--r-- 1 root root 3.2M Mar 17 05:28 pytorch_lora_weights.bin

We can push all these files nicely to the Hub thanks to huggingface_hub. The final repository on the Hub would look like this: https://huggingface.co/sayakpaul/fine-tuned-dreambooth/tree/main.

We can then do the loading part like so:

from peft import PeftConfig, PeftModel
from diffusers import DiffusionPipeline

# `model_id` can be determined programmatically from `repo_name`.
pipe = DiffusionPipeline.from_pretrained(model_id)

config = PeftConfig.from_pretrained(repo_name)
text_encoder = pipe.text_encoder
pipe.text_encoder = PeftModel.from_pretrained(text_encoder, repo_name)
pipe.unet.load_attn_procs(repo_name)
image = pipe("A picture of a dog in a bucket", num_inference_steps=10).images[0]

This is the Colab Notebook, where this above PoC can be found full-fledged.

One could argue that we can also use peft for the UNet part (it's possible). IIUC, in that case, we might have to maintain two separate repositories -- one for the UNet and another one for the text encoder. However, @pacman100 please correct me if I am wrong.

In any case, if the above (dual repo creation) issue is sorted out from peft, we can potentially deprecate our LoRA-related utilities.

If the above design looks good to you, I will drop a PR to modify the LoRA DreamBooth example.

Cc: @patrickvonplaten @williamberman

[patrickvonplaten]

I don't think we can have a new dependency just for LoRA text encoder especially for inference.
This complicates things too much. If necessary we should find a solution that purely resides in diffusers.

[sayakpaul]

I don't think we can have a new dependency just for LoRA text encoder especially for inference.

It's not just for inference. It's for training (with LoRA) too.

And we introduce this dependency only in the specific case where someone wants to do LoRA fine-tuning of the text encoder. I think the scope is in line with how we use wandb for validation inference.

This complicates things too much.

I am unable to see how. Could you elaborate or refer to me to any existing threads?

If necessary we should find a solution that purely resides in diffusers.

I think this, instead, would complicate things. But feel free to give your thoughts. I did consider it and couldn't think of a cleaner solution.

[patrickvonplaten]

peft is a training library not a model library. I don't think diffusers should rely on peft to run LoRA in model inference even as a soft-dependency. wandb is clearly just for training.

I'm mainly concerned about how to run LoRA models once they are trained. Imagine we have 500 trained LoRA checkpoins on the Hub that have been trained with peft, now we want to run those in inference. To run them in inference, we should not have to rely on peft.

To run diffusion models in inference, we should be able to just run:

from diffusers import DiffusionPipeline

# `model_id` can be determined programmatically from `repo_name`.
pipe = DiffusionPipeline.from_pretrained(model_id)

pipe.unet.load_attn_procs(repo_name)
image = pipe("A picture of a dog in a bucket", num_inference_steps=10).images[0]

(without having to install peft) instead of:

from peft import PeftConfig, PeftModel
from diffusers import DiffusionPipeline

# `model_id` can be determined programmatically from `repo_name`.
pipe = DiffusionPipeline.from_pretrained(model_id)

config = PeftConfig.from_pretrained(repo_name)
text_encoder = pipe.text_encoder
pipe.text_encoder = PeftModel.from_pretrained(text_encoder, repo_name)
pipe.unet.load_attn_procs(repo_name)
image = pipe("A picture of a dog in a bucket", num_inference_steps=10).images[0]

[patrickvonplaten]

Also note that peft relies on diffusers (which makes sense). It's weird to have diffusers rely on peft in return -> not a fan of circular dependencies like this

[sayakpaul]

So, for a LoRA fine-tuned text encoder, you'd ideally like:

from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(model_id)

pipe.unet.load_attn_procs(repo_name)
pipe.text_encoder.load_attn_procs(repo_name)
image = pipe("A picture of a dog in a bucket", num_inference_steps=10).images[0]

[patrickvonplaten]

Yeah exactly, the problem however is that the text_encoder is fully managed by transformers - so another idea could be to just add a:

from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(model_id)

pipe.load_lora(repo_name)

method directly to the DiffusionPipeline that calls unet.load_attn_procs(...) inside the function for the Unet LoRA layers and handles loading the text encoder layers into the DiffusionPipeline

[sayakpaul]

handles loading the text encoder layers into the DiffusionPipeline

Yeah, this should be possible but it doesn't consider the LoRA fine-tuned text encoder, which is what we're discussing here.

Maybe we should recommend following the peft training script for DreamBooth in case someone wants to LoRA fine-tune both (UNet and Text Encoder). This is because I don't see a clean (even semi-hacky) way around this thing. But happy to help with contributions if you have something in mind.

[sayakpaul]

@patrickvonplaten I thought about a non-peft solution.

My idea was to first initialize the LoRA layers for the text encoder like so:

from transformers import CLIPTextModel
from diffusers.models.attention_processor import LoRAAttnProcessor
from diffusers.loaders import AttnProcsLayers

text_encoder = CLIPTextModel.from_pretrained(
    "runwayml/stable-diffusion-v1-5", subfolder="text_encoder"
)

lora_attn_procs = {}

for name, module in text_encoder.named_modules():
    if "q_proj" in name or "v_proj" in name:
        lora_attn_procs[name] = LoRAAttnProcessor(
            hidden_size=module.out_features, cross_attention_dim=None
        )

text_encoder_lora_layers = AttnProcsLayers(lora_attn_procs)

Now, text_encoder_lora_layers can easily be trained and serialized just like how it's done for the UNet.

Now, after the LoRA layers have been trained, we need to perform "merging". For that, I came up with the following design:

rand_x = torch.randn(10, 768)

# Loop over the original attention modules.
for name, _ in text_encoder.named_modules():
    if "q_proj" in name or "v_proj" in name:
        # Retrieve the module and its corresponding LoRA processor.
        module = text_encoder.get_submodule(name)
        # Construct a new function that performs the LoRA merging. We will monkey patch 
        # this forward pass.
        new_forward = lambda x: module.forward(x) + lora_attn_procs[name].to_q_lora(x)
        
        # Some dummy forward pass for the current module and processor for quick testing.
        original_outputs = module.forward(rand_x)
        lora_outputs = new_forward(rand_x)

        # They shouldn't match.
        assert ~torch.allclose(original_outputs, lora_outputs)

        # Monkey-patch.
        module.forward = new_forward

I didn't consider the scale factor above but that can be easily added here I believe:

new_forward = lambda x: module.forward(x) + scale * lora_attn_procs[name].to_q_lora(x)

Here's my PoC Colab: https://colab.research.google.com/gist/sayakpaul/d372723a1723fce855e542442e95f6b7/scratchpad.ipynb

WDYT?

[patrickvonplaten]

I think that's a nice approach: #2719 (comment) @sayakpaul -> let's go for this :-)

[sayakpaul]

@patrickvonplaten I started working on it today. I found that after all the monkey-patching, the text encoder actually fails to do the forward pass.

To reproduce, consider the following minimal example.

Here's a sample class

import torch.nn as nn 

class GPT5(nn.Module):
    embed_dim = 768
    num_heads = 12
    q_proj = nn.Linear(embed_dim, embed_dim)
    head_dim = embed_dim // num_heads
    scale = head_dim**-0.5

    def forward(self, hidden_states):
        return self.q_proj(hidden_states) * self.scale

This works fine as expected:

gpt5 = GPT5()
gpt5(torch.randn(1, 10, 768)).size()

Do monkey-patching:

gpt5 = GPT5()
new_forward = lambda x: l.forward(x) + 1
gpt5.forward = new_forward

The following raises an error:

gpt5(torch.randn(1, 10, 768)).size()

─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ in <module>:1                                                                                    │
│                                                                                                  │
│ /usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py:1194 in _call_impl             │
│                                                                                                  │
│   1191 │   │   # this function, and just call forward.                                           │
│   1192 │   │   if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks o  │
│   1193 │   │   │   │   or _global_forward_hooks or _global_forward_pre_hooks):                   │
│ ❱ 1194 │   │   │   return forward_call(*input, **kwargs)                                         │
│   1195 │   │   # Do not call functions when jit is used                                          │
│   1196 │   │   full_backward_hooks, non_full_backward_hooks = [], []                             │
│   1197 │   │   if self._backward_hooks or _global_backward_hooks:                                │
│ in <lambda>:2                                                                                    │
│ in <lambda>:2                                                                                    │

Any pointers on approaching this?

[patrickvonplaten]

Ah interesting! Ok we need to maybe think about another solution maybe - will give it a thought this week