[LoRA] support kohya and xlabs loras for flux.

[sayakpaul]

What does this PR do?

Fixes #9291.
Fixes #9292 (potentially).

@apolinario does this work for you?

I decided to reuse https://github.com/kohya-ss/sd-scripts/blob/sd3/networks/convert_flux_lora.py (with appropriate credits to @kohya-ss) because the conversion utility won't have been too different if we were to write it anyway.

@kohya-ss could I get your commit email address so that we could add you as a co-author of this PR? @Leommm-byte same for you as well.

About tests, I think it might make sense to add a bit later to give the community a chance to play with it.

[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.

[apolinario]

The Kohya LoRA seem to be working!

import torch
from diffusers import DiffusionPipeline

device = "cuda" if torch.cuda.is_available() else "cpu"

pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
pipe = pipe.to(device)

pipe.load_lora_weights("Norod78/brain-slug-flux")
generator = torch.Generator(device="cpu").manual_seed(0)
prompt = "The girl with a brain slug earring"
image = pipe(prompt, generator=generator).images[0]
image

[kohya-ss]

@kohya-ss could I get your commit email address so that we could add you as a co-author of this PR? @Leommm-byte same for you as well.

Thank you for considering me as a co-author. I appreciate the credit, but there's no need to add me as a co-author to the PR. Please feel free to use the code as you see fit. I'm glad it could be helpful for your project.

[sayakpaul]

Sure, anyway, I have made a comment about the code I reused from sd-scripts. Thank you!

[apolinario]

XLabs LoRAs seem to be working, however they seem to need a higher scale to really shine/showcase their effect (also, using the guidance_scale of 3.5 instead of the default 7.0, as the super high guidance seems to outshine the effect of the LoRA)

A blue jay standing on a large basket of rainbow macarons, disney style (disney_lora.safetensors)
No LoRA	disney_lora.safetensors (scale 1.0)	disney_lora.safetensors (scale 1.75)

The portrait of a brazilian person (XLabs-AI/flux-RealismLora)
No LoRA	flux-RealismLora (scale 1.0)	flux-RealismLora (scale 1.75)

import torch
from diffusers import DiffusionPipeline

device = "cuda" if torch.cuda.is_available() else "cpu"

pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
pipe = pipe.to(device)

prompt = "A blue jay standing on a large basket of rainbow macarons, disney style"
#prompt = "The portrait of a brazilian person"

generator = torch.Generator(device="cpu").manual_seed(42)
image = pipe(prompt, generator=generator, guidance_scale=3.5).images[0]
image.save("no_lora.png")

pipe.load_lora_weights("XLabs-AI/flux-lora-collection", weight_name="disney_lora.safetensors")
#pipe.load_lora_weights("XLabs-AI/flux-RealismLora")

generator = torch.Generator(device="cpu").manual_seed(42)
image = pipe(prompt, generator=generator, joint_attention_kwargs={"scale": 1}, guidance_scale=3.5).images[0]
image.save("lora_1_0.png")

generator = torch.Generator(device="cpu").manual_seed(42)
image = pipe(prompt, generator=generator, joint_attention_kwargs={"scale": 1.75}, guidance_scale=3.5).images[0]
image.save("lora_1_75.png")

[apolinario]

One other thing, some LoRAs in this format have a diffusion_model prepend before the double_blocks, as in

diffusion_model.double_blocks.0.img_attn.proj.lora_down.weight

For example check this out

Probably worth to consider this possibility too in the conversion

[sayakpaul]

One other thing, some LoRAs in this format have a diffusion_model prepend before the double_blocks, as in

@apolinario thanks! Does d3c60e5 work for you?

[apolinario]

Let's go!

[Leommm-byte]

What does this PR do?

Fixes #9291. Fixes #9292 (potentially).

@apolinario does this work for you?

I decided to reuse https://github.com/kohya-ss/sd-scripts/blob/sd3/networks/convert_flux_lora.py (with appropriate credits to @kohya-ss) because the conversion utility won't have been too different if we were to write it anyway.

@kohya-ss could I get your commit email address so that we could add you as a co-author of this PR? @Leommm-byte same for you as well.

About tests, I think it might make sense to add a bit later to give the community a chance to play with it.

Ohh, I'm so glad I could be helpful with this challenge. I'm also honored to be considered a co-author.

If needed, here's my commit email address: leom20031@gmail.com

I'll test the adaptation as soon as I get the chance. Great work guys.

[yiyixuxu]

thanks!

[duonglegiang]

Do you have plans to support ControlNet and IP Adapter from XLabs-AI in future updates?
ControlNet: https://huggingface.co/XLabs-AI/flux-controlnet-canny
IP-Adapter: https://huggingface.co/XLabs-AI/flux-ip-adapter

[sayakpaul]

For ControlNet, #9301.

For IP-Adapters, cc: @fabiorigano @asomoza @yiyixuxu.

[Oguzhanercan]

High, thanks for your work. I wonder that, using trained loras-controlnets with 4-bit quantized DiTs(either nf4, raw int4 or SVD4Bit) is okey? Or should we quantize them also? As I see, flux-labs scripts are using float32 dtype. Loading and fusing them to 4bit dit is not throws an error. But I am not sure that loras trained with these scripts works with diffusers beacuse the output does not change at all.

@sayakpaul

[sayakpaul]

You need to keep some things in mind when fusing LoRAs into a quantized base model:
https://github.com/huggingface/diffusers/tree/main/examples/research_projects/flux_lora_quantization#inference

[Oguzhanercan]

You need to keep some things in mind when fusing LoRAs into a quantized base model:

https://github.com/huggingface/diffusers/tree/main/examples/research_projects/flux_lora_quantization#inference

Thank you

[dancasas]

I am using the provided code example to load a LoRA I trained with x-flux from X-Labs, but there seems to be some new keys that the current implementation of lora_conversion_utils.py does not parse. I get this error after pipe.load_lora_weights("/tmp/my_lora.safetensors")

ValueError: old_state_dict should be at this point but has: ['single_blocks.1.processor.qkv_lora.down.weight', 'single_blocks.1.processor.qkv_lora.up.weight', 'single_blocks.2.processor.qkv_lora.down.weight', 'single_blocks.2.processor.qkv_lora.up.weight', 'single_blocks.3.processor.qkv_lora.down.weight', 'single_blocks.3.processor.qkv_lora.up.weight', 'single_blocks.4.processor.qkv_lora.down.weight', 'single_blocks.4.processor.qkv_lora.up.weight'].

Does anyone know how to fix it? Thanks!

[sayakpaul]

Refer to #9915.

[dancasas]

Thanks! I can confirm that this PR #9915 worked for me to :)