[Flux LoRA] fix issues in flux lora scripts

[linoytsaban]

fix remaining pending issues from #10313, #9476 in Flux LoRA training scripts

• verify optimizer is updating properly (transformer only ☑️, test encoder w/ clip ☑️, pivotal w/ clip, pivotal w/ clip & t5, ti) wip
• accelerate error when running on multiple gpus
• replace scheduler
• log_validation with mixed precision
• save intermediate embeddings when checkpointing enabled

code snippets and output examples:

• for running log_validation with mixed precision-

import os
os.environ['MODEL_NAME'] = "black-forest-labs/FLUX.1-dev"
os.environ['DATASET_NAME'] = "dog"
os.environ['OUTPUT_DIR'] = "flux-test-1"

!accelerate launch train_dreambooth_lora_flux.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$DATASET_NAME \
  --output_dir=$OUTPUT_DIR \
  --mixed_precision="bf16" \
  --instance_prompt="a photo of sks dog" \
  --resolution=1024 \
  --train_batch_size=1 \
  --guidance_scale=1 \
  --gradient_accumulation_steps=1 \
  --optimizer="prodigy" \
  --learning_rate=1. \
  --report_to="wandb" \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=500 \
  --checkpointing_steps=250 \
  --validation_prompt="a photo of sks dog in a bucket"\
  --validation_epochs=25 \
  --seed="0" \
  --push_to_hub

validation output at step 380:

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

[luchaoqi]

Hi @linoytsaban , thanks for this prompt fix!

I believe the accelerator would produce error with line here with textual inversion specifically following the blog here:

[rank0]: Traceback (most recent call last):
[rank0]:   File "/playpen-nas-ssd/luchao/projects/diffusers/examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced_linoy.py", line 2576, in <module>
[rank0]:     main(args)
[rank0]:   File "/playpen-nas-ssd/luchao/projects/diffusers/examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced_linoy.py", line 2273, in main
[rank0]:     prompt_embeds, pooled_prompt_embeds, text_ids = encode_prompt(
[rank0]:                                                     ^^^^^^^^^^^^^^
[rank0]:   File "/playpen-nas-ssd/luchao/projects/diffusers/examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced_linoy.py", line 1446, in encode_prompt
[rank0]:     dtype = text_encoders[0].dtype
[rank0]:             ^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/playpen-nas-ssd/luchao/software/miniconda3/envs/diffuser/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1928, in __getattr__
[rank0]:     raise AttributeError(
[rank0]: AttributeError: 'DistributedDataParallel' object has no attribute 'dtype'. Did you mean: 'type'?

Also Is it possible to verify if textual inversion works in sks dog case on your end as well? e.g. pure CLIP textual inversion as mentioned here

  --train_text_encoder_ti \
  --train_text_encoder_ti_frac=1 \
  --train_transformer_frac=0

[linoytsaban]

hey @luchaoqi! yes I'm currently testing multiple configurations - will definitely test with pivotal tuning with clip and pure textual inversion with clip.

re: error with accelerator when running with multiple processes - adding it to the todo list :)

[sayakpaul]

Initial comments.

[linoytsaban]

@sayakpaul I noticed in the scripts some times we use: accelerator.unwrap_model, and sometimes we use unwrap_model

    def unwrap_model(model):
        model = accelerator.unwrap_model(model)
        model = model._orig_mod if is_compiled_module(model) else model
        return model

do you recall why it's not consistently one way or the other?

[sayakpaul]

Using the unwrap_model() function works. The ones that doesn't should be updated to have something similar. We added unwrap_model() to have more consistency for cases with torch.compile()

[linoytsaban]

Hey @luchaoqi! could you please check if the accelerator now works fine with distributed training? I think it should be resolved now

[luchaoqi]

Hi @linoytsaban, yes distributed training works as expected.

Pure textual inversion pops up new problems:

03/19/2025 10:20:03 - INFO - __main__ - Running validation...
 Generating 4 images with prompt: a photo of <s0><s1> person at 50 years old.
Traceback (most recent call last):
  File "/playpen-nas-ssd/luchao/projects/diffusers/examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced_linoy.py", line 2408, in <module>
    main(args)
  File "/playpen-nas-ssd/luchao/projects/diffusers/examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced_linoy.py", line 2055, in main
    text_encoder_one.train()
    ^^^^^^^^^^^^^^^^
UnboundLocalError: cannot access local variable 'text_encoder_one' where it is not associated with a value
[rank0]: Traceback (most recent call last):
[rank0]:   File "/playpen-nas-ssd/luchao/projects/diffusers/examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced_linoy.py", line 2408, in <module>
[rank0]:     main(args)
[rank0]:   File "/playpen-nas-ssd/luchao/projects/diffusers/examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced_linoy.py", line 2055, in main
[rank0]:     text_encoder_one.train()
[rank0]:     ^^^^^^^^^^^^^^^^
[rank0]: UnboundLocalError: cannot access local variable 'text_encoder_one' where it is not associated with a value

[luchaoqi]

Hi @linoytsaban, just wanted to follow up on the textual inversion part—do you anticipate it being fixed soon, or will it need a bit more time?

[linoytsaban]

@luchaoqi yes should be done soon!

[linoytsaban]

@luchaoqi if you want to give it a try the current version should be fixed

[luchaoqi]

@linoytsaban thanks! Would definitely try it out asap once I get some time.
Feel free to merge it if other reviewers agree, cheers!

[sayakpaul]

Thanks! Left some comments. Let me know if they make sense.

[sayakpaul]

I think this is only needed for the intermediate validation. Do we need to check for that?

[linoytsaban]

Yeah I think you're right, tested it now and seems to work as expected, changed it now

[sayakpaul]

Thanks a ton for handling this!

[sayakpaul]

@bot /style

[github-actions]

Style fixes have been applied. View the workflow run here.

[linoytsaban]

Failing test is unrelated