[LoRA] support Kohya Flux LoRAs that have text encoders as well

[sayakpaul]

What does this PR do?

https://huggingface.co/cocktailpeanut has a bunch of very nice Flux LoRAs that were trained using Kohya but has text encoder components too. This PR adds support for fully loading those LoRAs.

Test code (has a slow test in this PR too):

from diffusers import FluxPipeline
import torch 

pipeline = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16
).to("cuda")

pipeline.load_lora_weights("cocktailpeanut/optimus", weight_name="optimus.safetensors")
  
prompts = [
    "optimus is cleaning the house with broomstick",
    "optimus is a DJ performing at a hip nightclub",
    "optimus is competing in a bboy break dancing competition",
    "optimus is playing tennis in a tennis court"
]
images = pipeline(
    prompts, 
    num_inference_steps=50,
    guidance_scale=3.5,
    max_sequence_length=512,
    generator=torch.manual_seed(0)
).images
for i, image in enumerate(images):
    image.save(f"{i}.png")

Favorite sample:

^{optimus is a DJ performing at a hip nightclub}

[sayakpaul]

Sorry for the review request messup.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[yiyixuxu]

@asomoza can you give this a review too?

[sayakpaul]

Just for the sake of comparison, if we run the example code provided in the PR description with diffusers:main, we get:

with text encoder	without text encoder

[asomoza]

Thanks, LGTM in respect to the changes.

Since we now have text encoders training with a transformer model, I did some tests with the blockwise loras that I often use with SDXL:

no lora	just transformer	with TEs

This is not related to this PR but so that we just know.

if I do this:

scales = {"text_encoder": 0.0, "text_encoder_2": 0.0, "transformer": 0.0}
pipe.set_adapters("optimus", adapter_weights=scales)

It works as intended, but I copy & pasted the same code I use for SDXL:

scales = {"text_encoder": 0.0, "text_encoder_2": 0.0, "unet": 0.0}
pipe.set_adapters("optimus", adapter_weights=scales)

This doesn't work, the transformer keeps the lora scale at 1.0 but it doesn't show an error or warning that I'm setting the "unet" instead of the "transformer".

[sayakpaul]

@asomoza

scales = {"text_encoder": 0.0, "text_encoder_2": 0.0, "transformer": 0.0}

I think it shouldn't apply to this LoRA because it doesn't have the text_encoder_2 component in the first place. I can look into catching this and erroring/warning as needed. Will look into the "unet" thingy as well. Thanks much for flagging!

[apolinario]

Looking great! Thanks for adding it