The world’s leading publication for data science, AI, and ML professionals.

Hugging Face Diffusers can correctly load LoRA

Using the Latest Diffusers Monkey Patching function to load LoRA produces exactly the same result compare with A1111

Pull the latest code from Hugging Face’s Diffusers code repository, and found that the newest code updated related to LoRA loading is updated and can do Monkey-Patching LoRA loading now.

To install the latest Diffusers:

pip install -U git+https://github.com/huggingface/diffusers.git@main

The LoRA loading function was generating slightly faulty results yesterday, according to my test. This article discusses how to use the latest LoRA loader from the Diffusers package.

Load LoRA and update the Stable Diffusion model weight

It has been a while since programmers using Diffusers can’t have the LoRA loaded in an easy way. To load LoRA to a checkpoint model and output the same result as A1111’s Stable Diffusion Webui did, we need to use additional custom code to load the weights as I provided in this article.

Improving Diffusers Package for High-Quality Image Generation

The solution provided in this article works well and fast, while it requires additional management on the LoRA alpha weight, we need to create a variable to remember the current LoRA weight α. Because the load LoRA code simply adds put the A and B matrix from LoRA together.

Weight from LoRA
Weight from LoRA

And then merge with the main checkpoint model weight W.

Merge LoRA weight with checkpoint weight
Merge LoRA weight with checkpoint weight

To remove the LoRA weights, we will need a negative -α to remove the LoRA weights, or recreate the pipeline.

The Monkey-Patching way to load LoRA

Another way to use LoRA is patching the code that executes the module forward process, and bringing the LoRA weights during the time of calculating text embedding and attention score.

Monkey patching model with LoRA
Monkey patching model with LoRA

And this is how Diffusers LoraLoaderMixin’s approach to LoRA loading. The good part of this approach is that no model weight is updated, we can easily reset the LoRA and provide a new α to define the LoRA weight.

However, before today(July 26, 2023), Diffusers’ LoraLoaderMixin load LoRA and generate results somewhat different compared with A1111. And today’s code fixed the issue. "fixed" I mean, you can use Diffusers to load a checkpoint model together with LoRA and generate exactly the same result using A1111 SD webui.

Load LoRA using Diffusers LoraLoaderMixin

Assume we have the LoRA file in safetensor file format, to load LoRA using Diffusers is as easy as this.

import torch
from diffusers import StableDiffusionPipeline

text2img_pipe = StableDiffusionPipeline.from_pretrained(
    "stablediffusionapi/deliberate-v2"
    , torch_dtype = torch.float16
    , safety_checker = None
).to("cuda:0")

lora_path = "<path/to/lora.safetensors>"
text2img_pipe.load_lora_weights(lora_path)

Just one line code: text2img_pipe.load_lora_weights(lora_path). Test with one of the famous LoRA – LowRA , this LoRA can turn the image to dark mode.

Let’s test it with and without LoRA loaded.

from diffusers import EulerDiscreteScheduler

prompt = """
Маша making extreme selfie on skyscraper, bird's eye view, from above, night, smiling
"""
neg_prompt = """
NSFW,deformed, distorted, disfigured, poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, mutated hands and fingers, disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation
"""

text2img_pipe.scheduler = EulerDiscreteScheduler.from_config(text2img_pipe.scheduler.config)

image = text2img_pipe(
    prompt = prompt
    , negative_prompt = neg_prompt
    , generator = torch.Generator("cuda:0").manual_seed(3135098381)
    , num_inference_steps = 28
    , guidance_scale = 8
    , width = 512
    , height = 768
).images[0]
display(image)

Here are the result, the left is the one has no LoRA loaded, and the right one is the generated image with LoRA loaded using text2img_pipe.load_lora_weights(lora_path).

Figure 1. LowRA LoRA can generate dark mode image, generated by Andrew Zhu using Diffusers
Figure 1. LowRA LoRA can generate dark mode image, generated by Andrew Zhu using Diffusers

Load LoRA using Diffusers with a weight

Even though we can now load LoRA with just one line code, we still don’t see any support for the LoRA weight parameter. For example, we want less dark from the result and say add LoRA weight as 0.5. In A1111 Stable Diffusion Webui, we can give the LoRA weight as: <lora:LowRA:0.5>

How can we add this 0.5 to the Diffusers package? seems there is no way, but we can hack into the process and add our value in the middle of the LoRA loading like this:

text2img_pipe.unload_lora_weights()
lora_path = "<path/to/lora.safetensors>"

lora_w = 0.5
text2img_pipe._lora_scale = lora_w

state_dict, network_alphas = text2img_pipe.lora_state_dict(
    lora_path
)

for key in network_alphas:
    network_alphas[key] = network_alphas[key] * lora_w

#network_alpha = network_alpha * lora_w
text2img_pipe.load_lora_into_unet(
    state_dict = state_dict
    , network_alphas = network_alphas
    , unet = text2img_pipe.unet
)

text2img_pipe.load_lora_into_text_encoder(
    state_dict = state_dict
    , network_alphas = network_alphas
    , text_encoder = text2img_pipe.text_encoder
)

The left one is the result from A1111 SD Webui, and the right one from Diffusers, both using the same LoRA with 0.5 weight, Euler scheduler, and seed: 3135098381

Figure 2. Diffusers with LoRA can generate the same image compared with A1111. Image generated by Andrew Zhu using A1111 SD WebUi(left) and Diffusers(right)
Figure 2. Diffusers with LoRA can generate the same image compared with A1111. Image generated by Andrew Zhu using A1111 SD WebUi(left) and Diffusers(right)

Can you tell the difference?

To reset the LoRA, simply call the unload_lora_weights function :

text2img_pipe.unload_lora_weights()

Summary

This article discussed the most recent code update from Diffusers, an open-source Stable Diffusion package that can enable us to generate AI images using Python. Besides, we can freely add and change code according to specific needs.

For a long time, only A1111 Stable Diffusion WebUI can fully support LoRA, And now, we can use Python to control Diffusers to load thousands of community-shared LoRA and fully automate image generation, this opens a new door to GAI image generation.

In summary, the article’s contributions are as follows:

  1. Introduced the most simple way (one line code) to load LoRA safetensors file into Diffusers Stable Diffusion pipeline.
  2. Provide a simple hacky way to feed LoRA weights during the LoRA loading stage.

References

[1] Hugging Face, Low-Rank Adaptation of Large Language Models (LoRA), https://huggingface.co/docs/diffusers/training/lora

[2] LowRA, https://civitai.com/models/48139/lowra

[3] Load Kohya-ss style LoRAs with auxiliary states, https://github.com/huggingface/diffusers/pull/4147

[4] Simo Ryu(cloneofsimo), lora, https://github.com/cloneofsimo/lora


Related Articles