[Sd3 Dreambooth LoRA] Add text encoder training for the clip encoders

[linoytsaban]

add text encoder training support for the CLIP encoders to the dreambooth lora training script for SD3

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

[r-aristov]

When no individual prompts provided, only instance prompts, got this:

Traceback (most recent call last):
  File "/root/train_dreambooth_lora_sd3.py", line 1836, in <module>
    main(args)
  File "/root/train_dreambooth_lora_sd3.py", line 1631, in main
    encoder_hidden_states=prompt_embeds,
UnboundLocalError: local variable 'prompt_embeds' referenced before assignment

[linoytsaban]

thanks @r-aristov! I think it should be working now

[sayakpaul]

Can we confirm this via experiments that text encoder 3 training doesn't matter too much? Can be done separately and won't block this PR.

[apolinario]

Actually we skipped not because we think it doesn't matter but because we thought dealing with training the T5 would be a different animal than the already well known CLIP text encoder training (also on VRAM consumption side). So indeed we left it to a future PR to investigate the T5 training!

[sayakpaul]

Slightly worried about the dynamics about this so, let’s make sure we run ample experiments to see if training two text encoders while keeping the other one fixed works as expected.

[sayakpaul]

Thanks. I have left some comments. My main question is how much does training the text encoder matter here given we use three in SD3? Could we see some concrete comparative examples?

Additionally, we need to add tests to https://github.com/huggingface/diffusers/blob/main/tests/lora/test_lora_layers_sd3.py and add a note about --train_text_encoder in the REAMDE.

[linoytsaban]

some results for comparison

training config:

!accelerate launch train_dreambooth_lora_sd3.py \
  --pretrained_model_name_or_path="stabilityai/stable-diffusion-3-medium-diffusers"  \
  --dataset_name="Norod78/Yarn-art-style"\
  --output_dir="dreambooth-sd3-lora"\
  --mixed_precision="fp16" \
  --instance_prompt="a photo of TOK yarn art dog" \
  --resolution=1024 \
  --train_batch_size=1 \
  **--train_text_encoder\**
  --gradient_accumulation_steps=1 \
  --optimizer="prodigy"\
  --learning_rate=1.0 \
  **--text_encoder_lr=1.0\**
  --report_to="wandb" \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=1500 \
  --repeats=1\
  --rank=32\
  --weighting_scheme="logit_normal" \
  --validation_epochs=100 \
  --seed="0" \
  --push_to_hub

[sayakpaul]

Cool, the results are stunning. So, the TODOs are:

• Add relevant tests to https://github.com/huggingface/diffusers/blob/main/tests/lora/test_lora_layers_sd3.py
• Add documentation to the SD3 DreamBooth README
• Once [Tests] add test suite for SD3 DreamBooth #8650 is merged, I will take care of adding the text encoder training-related tests.

[sayakpaul]

Excellent work!