[Flux Dreambooth LoRA] - te bug fixes & updates

[linoytsaban]

• fixes to text encoder training bugs
• fix loss calculation
• fix param preparation - improve memory usage
• minor updates to readme
• add requirements
• add tests

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

[sayakpaul]

Let's add tests. No way otherwise, we can detect these. What say?

[Gothos]

training doesn't really seem to converge:

Granted, this is bsz 1, but at lr of 3e-6. This was supposed to be a person and only the outlines are visible.
Unless this is a lora loading issue? I think diffusers trained LoRA ought to be ok with diffusers loading?

[sayakpaul]

@Gothos could be because of some bug in the training, not sure if it's the loading because the lora_state_dict() and load_lora_weights() are from SD3, which we know works.

[Gothos]

@Gothos could be because of some bug in the training, not sure if it's the loading because the lora_state_dict() and load_lora_weights() are from SD3, which we know works.
Ok, thanks for the info!

[sayakpaul]

That was meant to be a piece of info, though. If you find anything dodgy holler at me.

[Gothos]

That was meant to be a piece of info, though. If you find anything dodgy holler at me.

Will do!

[linoytsaban]

I think it should be better now - with loss fix + memory usage is lower - @Gothos @arcanite24 if you want to give it a try
I did a 500 step training run on the dog dataset and it converged well -
prompt "an sks dog"

[Gothos]

Yep I can also confirm that it converges!

[sayakpaul]

Where did this go?

[linoytsaban]

this was removed because we discarded of precondition_outputs ,
originaly it was

# Follow: Section 5 of https://arxiv.org/abs/2206.00364.
# Preconditioning of the model outputs.
        if args.precondition_outputs:
            model_pred = model_pred * (-sigmas) + noisy_model_input

and was left accidentally in the previous merge after we removed precondition_outputs

[linoytsaban]

see also #9086 (comment)
for more context

[sayakpaul]

@linoytsaban let's maybe also include your additional results in the comments when they are available? I left one question on the changes and then we can merge I think.

[linoytsaban]

yarn art lora example (overfit a bit but nice vibe to it)

#!/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 Yarn art style tarot card"\
  --caption_column="text"\
  --output_dir="yarn_art_flux_500" \
  --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

[linoytsaban]

@sayakpaul let's merge once tests pass?

[sayakpaul]

Thank you!