DreamBooth training script / FLUX.1 [dev]

[lamm-mit]

Describe the bug

Thanks @bghira @sayakpaul for sharing the FLUX DreamBooth training scripts at https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/README_flux.md

I was trying to run the examples and also train on my own data (I have successfully used it to get fine-tunes using SD2.x, SDXL, etc.). While the LoRA training script works, I do not get a reduction in loss and am thinking that no training is happening (loss stays approximately constant, I tried a variety of LRs).

The full fine tune gives me a tensor dimension error:

[..] latents = latents.view(batch_size, height, width, channels // 4, 2, 2)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: shape '[1, 4, 4, 16, 2, 2]' is invalid for input of size 65536 [...]

The fine tune with text model -train_text_encoder also throws an error message. I was curious if others have tried it or if you know if someone could have another working example that I could try?

Thank you!

Also tagging @linoytsaban who I believe has done some pull requests/bug fixes

Reproduction

Examples in README: https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/README_flux.md

Logs

No response

System Info

diffusers-0.31.0.dev0

Who can help?

No response

[sayakpaul]

Cc: @linoytsaban

[linoytsaban]

hey @lamm-mit! is it possible you're not using the most recent version? merged in #9139

[lamm-mit]

Thanks for getting back to me @linoytsaban! Yes, I am using the latest merge (pip install from git and downloaded the training scripts directly from there also for my local directory). Any chance you have some loss plots or other data from your run so I can debug? Much appreciated.

[APannnnnnnnd]

Thanks for getting back to me @linoytsaban! Yes, I am using the latest merge (pip install from git and downloaded the training scripts directly from there also for my local directory). Any chance you have some loss plots or other data from your run so I can debug? Much appreciated.

So you solved the problem? I'm having the same problem. I haven't got a clue how to fix it yet.

[lamm-mit]

No, nothing yet. It still does not work. I had a lot of issues and in cases that did run, no training occurred (loss remained ~constant). Thank you everyone for your help to debug these issues.

[bghira]

loss remaining constant is a fixture of sigmoid timestep selection on a flow-matching model with non-augmented/preconditioned loss

[bghira]

here's a flat loss curve on a Flux model that trained really nicely - ignore the run names, i mislabeled them.

[APannnnnnnnd]

No, nothing yet. It still does not work. I had a lot of issues and in cases that did run, no training occurred (loss remained ~constant). Thank you everyone for your help to debug these issues.

model_pred = FluxPipeline._unpack_latents(
model_pred,
height=int(model_input.shape[2]) * 8,
width=int(model_input.shape[3]) * 8,
vae_scale_factor=vae_scale_factor,
)
#9139
I see that the code before the improvement had a * 8, and I didn't get an error after adding the height and width, but I'm not sure if this will cause other new problems.
So why did the source code delete this * 8 later?

[linoytsaban]

what resolution are you training with @APannnnnnnnd ?
for debugging purposes, could you try this config? (I used it for a lora that worked well)

!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

[APannnnnnnnd]

what resolution are you training with @APannnnnnnnd ? for debugging purposes, could you try this config? (I used it for a lora that worked well)

!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

Examples in README: https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/README_flux.md

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

accelerate launch train_dreambooth_flux.py
--pretrained_model_name_or_path="./FLUX.1-dev"
--instance_data_dir="mydata"
--output_dir="output_dir"
--mixed_precision="bf16"
--instance_prompt="a photo of sks dog"
--resolution=1024
--train_batch_size=1
--gradient_accumulation_steps=4
--learning_rate=1e-4
--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"
`

The example I used, it did have the above problem when training. After I added * 8 to the code there, it didn't give the error. Of course, I'm still debugging the code to fit my own project, but I don't think there will be any more problems.

[linoytsaban]

hmm I can't seem to reproduce your bug, currently training on 1024 resolution and it works fine so far
https://wandb.ai/linoy/dreambooth-flux-dev-lora/reports/loss-24-08-22-09-51-15---Vmlldzo5MTEyMTY5

#!/usr/bin/env bash
!accelerate launch train_dreambooth_lora_flux.py \
  --pretrained_model_name_or_path="black-forest-labs/FLUX.1-dev" \
  --dataset_name="Norod78/Yarn-art-style"\
  --instance_prompt="a puppy, yarn art style"\
  --output_dir="yarn_art_lora_flux" \
  --mixed_precision="bf16" \
  --use_8bit_adam\
  --weighting_scheme="none"\
  --resolution=1024 \
  --train_batch_size=1 \
  --repeats=1\
  --learning_rate=1e-4\
  --guidance_scale=1\
  --report_to="wandb"\
  --gradient_accumulation_steps=1 \
  --gradient_checkpointing \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --guidance_scale=1\
  --rank=4\
  --max_train_steps=700 \
  --checkpointing_steps=2000 \
  --seed="0"\
  --push_to_hub

for context, this loss is from another training I just ran, with same config, except different resolution (768)
didnt crash, and it did converge -

[APannnnnnnnd]

hmm I can't seem to reproduce your bug, currently training on 1024 resolution and it works fine so far https://wandb.ai/linoy/dreambooth-flux-dev-lora/reports/loss-24-08-22-09-51-15---Vmlldzo5MTEyMTY5

#!/usr/bin/env bash
!accelerate launch train_dreambooth_lora_flux.py \
  --pretrained_model_name_or_path="black-forest-labs/FLUX.1-dev" \
  --dataset_name="Norod78/Yarn-art-style"\
  --instance_prompt="a puppy, yarn art style"\
  --output_dir="yarn_art_lora_flux" \
  --mixed_precision="bf16" \
  --use_8bit_adam\
  --weighting_scheme="none"\
  --resolution=1024 \
  --train_batch_size=1 \
  --repeats=1\
  --learning_rate=1e-4\
  --guidance_scale=1\
  --report_to="wandb"\
  --gradient_accumulation_steps=1 \
  --gradient_checkpointing \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --guidance_scale=1\
  --rank=4\
  --max_train_steps=700 \
  --checkpointing_steps=2000 \
  --seed="0"\
  --push_to_hub

for context, this loss is from another training I just ran, with same config, except different resolution (768) didnt crash, and it did converge -

Oh, I see where the difference is, in train_dreambooth_lora_flux.py：
model_pred = FluxPipeline._unpack_latents( model_pred, height=int(model_input.shape[2] * vae_scale_factor / 2), width=int(model_input.shape[3] * vae_scale_factor / 2), vae_scale_factor=vae_scale_factor, )
The value of vae_scale_factor should be 16 by default.
But in the file train_dreambooth_flux.py that I'm running,
model_pred = FluxPipeline._unpack_latents( model_pred, height=int(model_input.shape[2]), width=int(model_input.shape[3]), vae_scale_factor=vae_scale_factor, )
There is no * 8 in the source code. After I added * 8 to height and width, there was no error.

Still, thanks for your help and advice.

[linoytsaban]

oh ok I see, let me look into the dreambooth script as well!