the dog example of train_dreambooth_lora_flux.py can not convergence

[chongxian]

Describe the bug

export MODEL_NAME="black-forest-labs/FLUX.1-dev"
export INSTANCE_DIR="dog"
export OUTPUT_DIR="trained-flux-lora"

accelerate launch train_dreambooth_lora_flux.py \
 --pretrained_model_name_or_path=$MODEL_NAME \
 --instance_data_dir=$INSTANCE_DIR \
 --output_dir=$OUTPUT_DIR \
 --mixed_precision="bf16" \
 --instance_prompt="a photo of sks dog" \
 --resolution=512 \
 --train_batch_size=1 \
 --gradient_accumulation_steps=4 \
 --learning_rate=1e-5 \
 --report_to="wandb" \
 --lr_scheduler="constant" \
 --lr_warmup_steps=0 \
 --max_train_steps=500 \
 --validation_prompt="A photo of sks dog in a bucket" \
 --validation_epochs=25 \
 --seed="0" \
 --push_to_hub

I follow this command to train lora of flux-dev and download the dog-example from huggingFace, but this setting could not get better result, the loss is normal

the dog-example look like this:

but my result look like below:

and don't use the lora to generate image of the same prompt look like below:

Reproduction

import torch
from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained("/opt/ml/volume/default/aigc/project/FLUX.1-dev",torch_dtype=torch.bfloat16)
pipe.enable_model_cpu_offload()
pipe.lora_state_dict("/opt/ml/volume/default/aigc/project/diffusers/examples/dreambooth/trained-flux-lora/checkpoint-500")
prompts = []
prompts.append("an sks dog")
index = 0
for prompt in prompts:
 image = pipe(
 prompt=prompt,
 num_inference_steps=20,
 guidance_scale=7.5,
 max_sequence_length=512,
 width=1152,
 height=768
 ).images[0]
 save_file = "dog"+str(index)+'.png'
 index+=1
 image.save(save_file)

Logs

No response

System Info

ubuntu 20.04

Who can help?

@sayakpaul @linoytsaban

[linoytsaban]

Hey @chongxian, did you try the most recent version (from this #9139 merged PR)? it should converge well :)

[chongxian]

Hey @chongxian, did you try the most recent version (from this #9139 (comment) merged PR)? it should converge well :)

I use the latest version, but I can't get the result like your result

[linoytsaban]

I think it could be the learning rate here,
could you try the setting I used?

!accelerate launch train_dreambooth_lora_flux.py \
 --pretrained_model_name_or_path="black-forest-labs/FLUX.1-dev" \
 --dataset_name="dog"\
 --instance_prompt="an sks dog"\
 --output_dir="dog_lora_flux" \
 --mixed_precision="bf16" \
 --optimizer="prodigy" \
 --weighting_scheme="none"\
 --resolution=512 \
 --train_batch_size=1 \
 --repeats=1\
 --learning_rate=1.0\
 --report_to="wandb"\
 --gradient_accumulation_steps=1 \
 --gradient_checkpointing \
 --lr_scheduler="constant" \
 --lr_warmup_steps=0 \
 --rank=4\
 --max_train_steps=500 \
 --checkpointing_steps=2000 \
 --seed="0"\
 --push_to_hub

[chongxian]

I think it could be the learning rate here, could you try the setting I used?

!accelerate launch train_dreambooth_lora_flux.py \
 --pretrained_model_name_or_path="black-forest-labs/FLUX.1-dev" \
 --dataset_name="dog"\
 --instance_prompt="an sks dog"\
 --output_dir="dog_lora_flux" \
 --mixed_precision="bf16" \
 --optimizer="prodigy" \
 --weighting_scheme="none"\
 --resolution=512 \
 --train_batch_size=1 \
 --repeats=1\
 --learning_rate=1.0\
 --report_to="wandb"\
 --gradient_accumulation_steps=1 \
 --gradient_checkpointing \
 --lr_scheduler="constant" \
 --lr_warmup_steps=0 \
 --rank=4\
 --max_train_steps=500 \
 --checkpointing_steps=2000 \
 --seed="0"\
 --push_to_hub

thanks,I'll try this setting

[chongxian]

I think it could be the learning rate here, could you try the setting I used?

!accelerate launch train_dreambooth_lora_flux.py \
 --pretrained_model_name_or_path="black-forest-labs/FLUX.1-dev" \
 --dataset_name="dog"\
 --instance_prompt="an sks dog"\
 --output_dir="dog_lora_flux" \
 --mixed_precision="bf16" \
 --optimizer="prodigy" \
 --weighting_scheme="none"\
 --resolution=512 \
 --train_batch_size=1 \
 --repeats=1\
 --learning_rate=1.0\
 --report_to="wandb"\
 --gradient_accumulation_steps=1 \
 --gradient_checkpointing \
 --lr_scheduler="constant" \
 --lr_warmup_steps=0 \
 --rank=4\
 --max_train_steps=500 \
 --checkpointing_steps=2000 \
 --seed="0"\
 --push_to_hub

I use this setting like you, just modify the datset_name to insrance_data_dir, remove the wandb and checkpointing_steps=2000

export MODEL_NAME="/opt/ml/volume/default/aigc/project/FLUX.1-dev"
export INSTANCE_DIR="/opt/ml/volume/default/aigc/project/diffusers/examples/dreambooth/dog"
export OUTPUT_DIR="./dog_lora_flux"


accelerate launch train_dreambooth_lora_flux.py \
 --pretrained_model_name_or_path=$MODEL_NAME \
 --instance_data_dir=$INSTANCE_DIR \
 --output_dir=$OUTPUT_DIR \
 --instance_prompt="an sks dog" \
 --mixed_precision="bf16" \
 --optimizer="prodigy" \
 --weighting_scheme="none" \
 --resolution=512 \
 --train_batch_size=1 \
 --repeats=1 \
 --learning_rate=1.0 \
 --gradient_accumulation_steps=1 \
 --gradient_checkpointing \
 --lr_scheduler="constant" \
 --lr_warmup_steps=0 \
 --rank=4 \
 --max_train_steps=500 \
 --seed="0"

however,the result is also bad

import torch
from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained("/opt/ml/volume/default/aigc/project/FLUX.1-dev",torch_dtype=torch.bfloat16)
pipe.enable_model_cpu_offload()
pipe.lora_state_dict("/opt/ml/volume/default/aigc/project/diffusers/examples/dreambooth/dog_lora_flux/checkpoint-500")
prompts = []
prompts.append("an sks dog")
# prompts.append("a man and a woman standing in the sand and look up ,dutch angle")
# prompts.append("a couple of people standing on a set of stairs ,dutch angle")
# prompts.append("a man holding a glass of wine next to another man sitting on a couch,dutch angle")
index = 0
for prompt in prompts:
 image = pipe(
 prompt=prompt,
 num_inference_steps=20,
 guidance_scale=7.5,
 max_sequence_length=512,
 width=1152,
 height=768
 ).images[0]
 save_file = "dog"+str(index)+'.png'
 index+=1
 image.save(save_file)

[linoytsaban]

from your snippet it appears you're simply not loading the lora :)

it should be pipe.load_lora_weights("path_to_your_lora") instead of pipe.lora_state_dict()

could you try again?

[chongxian]

from your snippet it appears you're simply not loading the lora :)

it should be pipe.load_lora_weights("path_to_your_lora") instead of pipe.lora_state_dict()

could you try again?

thanks for your reply, it work