train_dreambooth_lora_flux_advanced ValueError: unexpected save model: <class 'transformers.models.t5.modeling_t5.T5EncoderModel'>

[hackkhai]

Describe the bug

Was exploring train_dreambooth_lora_flux_advanced.py from the advanced diffusion training examples. When I used --enable_t5_ti was getting the error ValueError: unexpected save model: <class 'transformers.models.t5.modeling_t5.T5EncoderModel'>

Reproduction

export MODEL_NAME="black-forest-labs/FLUX.1-dev"
export DATASET_NAME=""
export INSTANCE_DIR="images"
export OUTPUT_DIR="test-LoRA"
accelerate launch train_dreambooth_lora_flux_advanced.py
--pretrained_model_name_or_path=$MODEL_NAME
--instance_data_dir=$INSTANCE_DIR
--instance_prompt="photo of TOK"
--output_dir=$OUTPUT_DIR
--mixed_precision="bf16"
--resolution=1024
--train_batch_size=1
--repeats=1
--report_to="tensorboard"
--gradient_accumulation_steps=1
--gradient_checkpointing
--learning_rate=1.0
--optimizer="prodigy"
--train_text_encoder_ti
--train_text_encoder_ti_frac=0.5
--lr_scheduler="constant"
--lr_warmup_steps=0
--rank=16
--max_train_steps=1000
--checkpointing_steps=100
--seed="42"
--text_encoder_lr=1.0
--enable_t5_ti \

Logs

Traceback (most recent call last):
  File "/home/azureuser/dreambooth_experiments/diffusers_flux/diffusers/examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced.py", line 2463, in <module>
    main(args)
  File "/home/azureuser/dreambooth_experiments/diffusers_flux/diffusers/examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced.py", line 2346, in main
    accelerator.save_state(save_path)
  File "/anaconda/envs/diff_test/lib/python3.10/site-packages/accelerate/accelerator.py", line 3081, in save_state
    hook(self._models, weights, output_dir)
  File "/home/azureuser/dreambooth_experiments/diffusers_flux/diffusers/examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced.py", line 1654, in save_model_hook
    raise ValueError(f"unexpected save model: {model.__class__}")
ValueError: unexpected save model: <class 'transformers.models.t5.modeling_t5.T5EncoderModel'>

System Info

A100 PCIE

Who can help?

@yiyixuxu

[a-r-r-o-w]

cc @linoytsaban for advanced training scripts

[linoytsaban]

Thanks to @rootonchair's PR, this should be fixe now🙌🏻

[Nomination-NRB]

Thanks to @rootonchair's PR, this should be fixe now🙌🏻

Could you tell me how to install diffusers version=0.32.0.dev0 ?
I use pip install diffusers and get version=0.31.0

[rootonchair]

@Nomination-NRB you should clone the repo then pip install -e . to get the dev version

[Nomination-NRB]

@Nomination-NRB you should clone the repo then pip install -e . to get the dev version

Thanks you so much, I successfully install the 0.32.0.dev0 according to your advice.