Issue with Flux LoRAs trained with SimpleTuner

[sayakpaul]

@bghira we're seeing issues when doing:

import torch
from diffusers import DiffusionPipeline

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

pipe.load_lora_weights("pzc163/flux-lora-littletinies")
generator = torch.Generator(device="cuda").manual_seed(0)
image = pipe(
    prompt=f"a dog in the park",
    num_inference_steps=28,
    guidance_scale=3.5,
    width=1024,
    height=1024,
    generator=generator,
 ).images[0]

But it outputs gibberish. The inference code from https://huggingface.co/pzc163/flux-lora-littletinies shows many unsupported arguments such negative_prompt and guidance_rescale.

Have you found any bugs with inference?

Cc: @apolinario @asomoza

[AmericanPresidentJimmyCarter]

+1, no idea how to use it with the quantized model either.

[sayakpaul]

@apolinario says Comfy can run it perfectly fine. So, I wonder what we're missing.

[sayakpaul]

Okay, it's the lora alpha that is the culprit

diff --git a/src/diffusers/loaders/lora_pipeline.py b/src/diffusers/loaders/lora_pipeline.py
index f612cc0c6..f432175d0 100644
--- a/src/diffusers/loaders/lora_pipeline.py
+++ b/src/diffusers/loaders/lora_pipeline.py
@@ -1686,6 +1686,7 @@ class FluxLoraLoaderMixin(LoraBaseMixin):
                     )
                 else:
                     lora_config_kwargs.pop("use_dora")
+            lora_config_kwargs['lora_alpha'] = 16
             lora_config = LoraConfig(**lora_config_kwargs)
 
             # adapter_name

With that hack and the following code,

import torch
from diffusers import DiffusionPipeline

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

pipe.load_lora_weights("pzc163/flux-lora-littletinies")
generator = torch.Generator(device="cuda").manual_seed(0)
image = pipe(
    prompt=f"ethnographic photography of teddy bear at a picnic",
    num_inference_steps=50,
    guidance_scale=7.5,
    max_sequence_length=512,
    width=1152,
    height=768,
    generator=generator,
).images[0]
image.save("flux_dev_lora.png")

I am able to generate:

[sayakpaul]

I also see SimpleTuner allows to train a different lora_alpha than lora_rank. For diffusers, it's always the same rank for easier loading:

https://github.com/bghira/SimpleTuner/blob/513b71b65b85af3bfc6c8ce1ee8abaafef1f4152/train.py#L932C28-L932C43

I wonder how lora_alpha is parsed for diffusers LoRA params then. Because we parse alphas only for non-diffusers checkpoints

diffusers/src/diffusers/loaders/lora_pipeline.py

Line 225 in 65e3090

state_dict, network_alphas = _convert_non_diffusers_lora_to_diffusers(state_dict)

The problem is a little bit tricky because now we have to either do some kind of heuristics to determine the alpha or allow saving LoRA configs in the metadata of the LoRA checkpoint. This is only possible if the checkpoint is serialized in safetensors, which should not be a problem.

[sayakpaul]

We cannot likely obtain the lora_alpha info from the input state dict. To verify, I did the following:

from diffusers import FluxTransformer2DModel 
from peft import LoraConfig
from peft.utils import get_peft_model_state_dict

init_dict = {
    "patch_size": 1,
    "in_channels": 4,
    "num_layers": 1,
    "num_single_layers": 1,
    "attention_head_dim": 16,
    "num_attention_heads": 2,
    "joint_attention_dim": 32,
    "pooled_projection_dim": 32,
    "axes_dims_rope": [4, 4, 8],
} 
model = FluxTransformer2DModel(**init_dict)

denoiser_lora_config = LoraConfig(
    r=4,
    lora_alpha=2,
    target_modules=["to_q", "to_k", "to_v", "to_out.0"],
    init_lora_weights=False
)
model.add_adapter(denoiser_lora_config)
peft_sd = get_peft_model_state_dict(model)
for k, v in peft_sd.items():
    print(k ,v.shape)

Despite the lora_alpha being set to 2 (different from rank), it's nowhere to be seen in the state dict shapes and that is expected:

transformer_blocks.0.attn.to_q.lora_A.weight torch.Size([4, 32])
transformer_blocks.0.attn.to_q.lora_B.weight torch.Size([32, 4])
transformer_blocks.0.attn.to_k.lora_A.weight torch.Size([4, 32])
transformer_blocks.0.attn.to_k.lora_B.weight torch.Size([32, 4])
transformer_blocks.0.attn.to_v.lora_A.weight torch.Size([4, 32])
transformer_blocks.0.attn.to_v.lora_B.weight torch.Size([32, 4])
transformer_blocks.0.attn.to_out.0.lora_A.weight torch.Size([4, 32])
transformer_blocks.0.attn.to_out.0.lora_B.weight torch.Size([32, 4])
single_transformer_blocks.0.attn.to_q.lora_A.weight torch.Size([4, 32])
single_transformer_blocks.0.attn.to_q.lora_B.weight torch.Size([32, 4])
single_transformer_blocks.0.attn.to_k.lora_A.weight torch.Size([4, 32])
single_transformer_blocks.0.attn.to_k.lora_B.weight torch.Size([32, 4])
single_transformer_blocks.0.attn.to_v.lora_A.weight torch.Size([4, 32])
single_transformer_blocks.0.attn.to_v.lora_B.weight torch.Size([32, 4])

So, we probably have to save the LoRA config in the safetensors metadata as I explained here:
#9134 (comment). I did a PoC for this back in the day: #6135.

@BenjaminBossan do you have any other idea here?

[sayakpaul]

Okay, another interesting fact discovered by @asomoza is that this LoRA should work as expected with schnell and with dev the results are widely different.

[sayakpaul]

@AmericanPresidentJimmyCarter I did some brute-forcing (it's not long-term) to perform inference with quantization:

import torch
from diffusers import DiffusionPipeline
from optimum.quanto import quantize, qfloat8, freeze

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

pipe.load_lora_weights("pzc163/flux-lora-littletinies")
# generator = torch.Generator(device="cuda").manual_seed(0)

pipe.fuse_lora(lora_scale=0.75)
pipe.unload_lora_weights()

quantize(pipe.transformer, qfloat8)
freeze(pipe.transformer)

# quantize(pipe.text_encoder, qfloat8)
# freeze(pipe.text_encoder)

# quantize(pipe.text_encoder_2, qfloat8)
# freeze(pipe.text_encoder_2)

image = pipe(
    prompt=f"a girl is walking in the forest,meeting a big tiger",
    num_inference_steps=4,
    guidance_scale=0.0,
    max_sequence_length=256,
    width=1024,
    height=1024,
    # joint_attention_kwargs={"scale": 1.0},
    generator=torch.manual_seed(0),
).images[0]
image.save("flux_dev_lora.png")

But it still doesn't apparently match the ComfyUI output even with lora_alpha=16 which is the default alpha value SimpleTuner sets:
https://github.com/bghira/SimpleTuner/blob/aceaf3e0b69de83f5a5864a5c40676982f4497cc/helpers/arguments.py#L240

Interestingly, if we keep the lora_alpha between [32, 42], the results somewhat resemble ComfyUI but not entirely.

[bghira]

i also received a report that lora alpha is in the state dict for other loras eg kohya, which makes it hard to know how to load diffusers loras which do not

[bghira]

i wonder if it needs to be scaled in fp32 and not bf16

[sayakpaul]

Yeah for other LoRAs, we parse the alpha but for diffusers we don't. This is why in our training code, lora_alpha=args.rank always. To allow, alpha, we would need something like this: #6135.

We never serialize alpha information in our LoRA ckpts and this is because historically, we never supported them from our training.

i wonder if it needs to be scaled in fp32 and not bf16

You mean the LoRA layers? Despite that, we need to change lora_alpha here to 16 (I am assuming 16 was used to train the said LoRA as that's default in SimpleTuner).

diffusers/src/diffusers/loaders/lora_pipeline.py

Line 1689 in 65e3090

lora_config = LoraConfig(**lora_config_kwargs)

Still trying to identify the potential suspects.

[bghira]

the rank and alpha should be / is listed in the model card too fwiw

[sayakpaul]

Yes, it is listed. Rank is 64 and alpha is 16.

[BenjaminBossan]

Hmm, apart from storing the alpha value somewhere in the metadata to access it later, there is one thing that comes to mind: In the end, alpha is just a scale for the LoRA output: output = base_output + lora_output * alpha / rank. IIUC we assume by default that alpha = rank. Let's say we have an example where alpha = 2 * rank, the lora_output would be off by a factor of 2. Therefore, we could achieve the same results by multiplying 2 onto one of the LoRA weights, say lora_A.

To make this work would require an export script that is aware of alpha and knows that the output is intended to be loaded into diffusers. Then it could go through the state_dict and multiply each lora_A by alpha / rank.

This is not ideal and I would prefer if alpha was known through the metadata, but maybe it's an idea.

[sayakpaul]

Thanks for providing your suggestion. However, alpha should be determined during loading just like we determine the rank. IIUC, to follow your suggestion, we would need to apply that scaling during serializing, right?

[bghira]

is there anything stopping us from adding alpha to the safetensors data? no one uses pth files anymore but even if they did it would probably be doable there too