Dreambooth LoRA Flux training last step error

[PluginBOXone]

Describe the bug

As soon as the training is done and the code wants to clear up and do its last steps I get this error

Steps: 99%|█████████▉| 397/400 [07:45<00:03, 1.16s/it, loss=0.397, lr=1]
Steps: 100%|█████████▉| 398/400 [07:47<00:02, 1.20s/it, loss=0.397, lr=1]
Steps: 100%|█████████▉| 398/400 [07:47<00:02, 1.20s/it, loss=0.539, lr=1]
Steps: 100%|█████████▉| 399/400 [07:48<00:01, 1.19s/it, loss=0.539, lr=1]
Steps: 100%|█████████▉| 399/400 [07:48<00:01, 1.19s/it, loss=0.58, lr=1]
Steps: 100%|██████████| 400/400 [07:49<00:00, 1.18s/it, loss=0.58, lr=1]
Steps: 100%|██████████| 400/400 [07:49<00:00, 1.18s/it, loss=0.288, lr=1] Model weights saved in /workspace/output_model/dd304483-afdc-4398-9c46-c660d0725e70-e1/pytorch_lora_weights.safetensors
2025-02-19T21:38:48.866518894Z Traceback (most recent call last):
2025-02-19T21:38:48.866557263Z File "/workspace/./train_dreambooth_lora_flux.py", line 1935, in
2025-02-19T21:38:48.867054758Z main(args)
2025-02-19T21:38:48.867072609Z File "/workspace/./train_dreambooth_lora_flux.py", line 1887, in main
2025-02-19T21:38:48.867457265Z pipeline = FluxPipeline.from_pretrained(
2025-02-19T21:38:48.867479814Z ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-02-19T21:38:48.867487574Z File "/usr/local/lib/python3.11/dist-packages/huggingface_hub/utils/_validators.py", line 114, in _inner_fn
2025-02-19T21:38:48.867554504Z return fn(*args, **kwargs)
2025-02-19T21:38:48.867601603Z ^^^^^^^^^^^^^^^^^^^
2025-02-19T21:38:48.867606703Z File "/usr/local/lib/python3.11/dist-packages/diffusers/pipelines/pipeline_utils.py", line 793, in from_pretrained
2025-02-19T21:38:48.867905410Z config_dict = cls.load_config(cached_folder, dduf_entries=dduf_entries)
2025-02-19T21:38:48.867973610Z ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-02-19T21:38:48.867992349Z File "/usr/local/lib/python3.11/dist-packages/huggingface_hub/utils/_validators.py", line 114, in _inner_fn
2025-02-19T21:38:48.868030829Z return fn(*args, **kwargs)
2025-02-19T21:38:48.868039849Z ^^^^^^^^^^^^^^^^^^^
2025-02-19T21:38:48.868053699Z File "/usr/local/lib/python3.11/dist-packages/diffusers/configuration_utils.py", line 381, in load_config
2025-02-19T21:38:48.868183318Z raise EnvironmentError(
2025-02-19T21:38:48.868199778Z OSError: Error no file named model_index.json found in directory /workspace/model/realflux1.
2025-02-19T21:38:49.009733209Z
Steps: 100%|██████████| 400/400 [07:49<00:00, 1.17s/it, loss=0.288, lr=1]
2025-02-19T21:38:50.343330576Z Traceback (most recent call last):
2025-02-19T21:38:50.343379125Z File "/usr/local/bin/accelerate", line 8, in
2025-02-19T21:38:50.343568443Z sys.exit(main())
2025-02-19T21:38:50.343608783Z ^^^^^^
2025-02-19T21:38:50.343691572Z File "/usr/local/lib/python3.11/dist-packages/accelerate/commands/accelerate_cli.py", line 48, in main
2025-02-19T21:38:50.343770471Z args.func(args)
2025-02-19T21:38:50.343896080Z File "/usr/local/lib/python3.11/dist-packages/accelerate/commands/launch.py", line 1106, in launch_command
2025-02-19T21:38:50.344200407Z simple_launcher(args)
2025-02-19T21:38:50.344262447Z File "/usr/local/lib/python3.11/dist-packages/accelerate/commands/launch.py", line 704, in simple_launcher
2025-02-19T21:38:50.344583144Z raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
2025-02-19T21:38:50.819696870Z ✅ FLUX LoRA Training abgeschlossen!

So basically the error says OSError: Error no file named model_index.json found in directory /workspace/model/realflux1. File "/workspace/./train_dreambooth_lora_flux.py", line 1887, in main
2025-02-19T21:38:48.867457265Z pipeline = FluxPipeline.from_pretrained(

But the path exists and the model in it. It is the same path that I started the training with it and it found the files/works.

It still looks like it finished the training because it saved a 98MB big .safetensors file and a log... but i have the feeling the LoRA is broken because when I try to load it, the inference output ist corrupted.

Tried inference without loading the LoRA:

As soon as loading the LoRA (i even tried different prompts that dont have todo anything with the LoRA)

Reproduction

Dreambooth Training:
Using the newest train_dreambooth_lora_flux.py, started with these parameters (of course with accelerate): ./train_dreambooth_lora_flux.py', '--pretrained_model_name_or_path', '/workspace/model/realflux1', '--instance_data_dir', '/workspace/job_files/dd304483-afdc-4398-9c46-e1/clean_data', '--output_dir', '/workspace/output_model/dd304483-afdc-4398-9c46-e1', '--instance_prompt', 'photo of WIXBSAHA black car', '--resolution', '768', '--learning_rate', '1.0', '--mixed_precision', 'bf16', '--lr_warmup_steps', '0', '--gradient_accumulation_steps', '1', '--lr_scheduler', 'constant', '--train_batch_size', '1', '--max_train_steps', '400', '--checkpointing_steps', '500', '--num_train_epochs', '10', '--checkpoints_total_limit', '1', '--train_text_encoder', '--rank', '16', '--optimizer', 'prodigy', '--repeats', '3', '--guidance_scale', '1'

Inference:

print("Lade FLUX Modell")
    pipe = FluxPipeline.from_pretrained(model_path, torch_dtype=torch.bfloat16).to("cuda")
    pipe.enable_model_cpu_offload()

    generator = None

    if lora_path:
        print(f"🔄 Lade FLUX LoRA-Modell: {lora_path}")
        pipe.load_lora_weights(lora_path)
        print("✅ LoRA geladen.")
    
    if seed is not None:
        generator = torch.Generator(device="cpu").manual_seed(seed)

    image = pipe(
        prompt=prompt,
        guidance_scale=guidance_scale, #0.
        negative_prompt=negative_prompt,
        height=height,
        true_cfg_scale=true_cfg_scale,
        width=width,
        num_inference_steps=num_inference_steps,
        max_sequence_length=max_sequence_length, #256
        generator=generator
    ).images[0]

image.save("test.png")

I also get a warning on inference:

✅ LoRA geladen.\n
[info]🔄 Lade FLUX LoRA-Modell: /workspace/lora_output_model\n
[info]\rLoading pipeline components...: 57%|█████▋ | 4/7 [00:00<00:00, 4.14it/s]\rLoading pipeline components...: 71%|███████▏ | 5/7 [00:01<00:00, 2.34it/s]\rLoading pipeline components...: 100%|██████████| 7/7 [00:01<00:00, 4.16it/s]\n
[info]\rLoading checkpoint shards: 100%|██████████| 2/2 [00:00<00:00, 7.39it/s]�[A\rLoading checkpoint shards: 100%|██████████| 2/2 [00:00<00:00, 7.37it/s]\n
[info]\rLoading checkpoint shards: 50%|█████ | 1/2 [00:00<00:00, 7.33it/s]�[A\n
[info]\rLoading checkpoint shards: 0%| | 0/2 [00:00<?, ?it/s]�[A\n
[info]\rLoading pipeline components...: 43%|████▎ | 3/7 [00:00<00:00, 5.06it/s]\n
[info]\rLoading pipeline components...: 0%| | 0/7 [00:00<?, ?it/s]\rLoading pipeline components...: 14%|█▍ | 1/7 [00:00<00:00, 7.55it/s]\rLoading pipeline components...: 29%|██▊ | 2/7 [00:00<00:00, 5.91it/s]You set add_prefix_space. The tokenizer needs to be converted from the slow tokenizers\n
[info]Lade FLUX Modell\n
[info]FLUX Inference\n

(read from bottom to top)

this one: You set add_prefix_space. The tokenizer needs to be converted from the slow tokenizers\n

Logs

System Info

A100 80GB VRAM 120GB RAM
pytorch:2.4.0-py3.11-cuda12.4.1
CUDA 12.4
accelerate==0.33.0

Who can help?

@sayakpaul

[sayakpaul]

Can you properly format the logs? It's very hard to parse through it.

[PluginBOXone]

Can you properly format the logs? It's very hard to parse through it.

Is it better now? 😬

[sayakpaul]

What happens when you try one of the example training commands from
https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/README_flux.md?

If it goes successfully, then it's probably something wrong on your end.

[PluginBOXone]

Interesting, I tried the exact parameters from the README Dreambooth + LoRA. It is now crashing at the beginning with the same error. It cant find the model although it can find the model on loading and training??? Really weird. I get these two warnings at the beginning but I dont think they matter.

--TRAIN_DREAMBOOTH_LORA_FLUX.py RUNNING--
2025-02-20 12:23:12,506 - INFO - Distributed environment: DistributedType.NO
Num processes: 1
Process index: 0
Local process index: 0
Device: cuda
Mixed precision type: bf16
You set add_prefix_space. The tokenizer needs to be converted from the slow tokenizers
You are using a model of type clip_text_model to instantiate a model of type . This is not supported for all configurations of models and can yield errors.
You are using a model of type t5 to instantiate a model of type . This is not supported for all configurations of models and can yield errors.

And this is the error I get as soon as it is training and trying to create validation Images:

2025-02-20 12:25:04,112 - INFO - Total train batch size (w. parallel, distributed & accumulation) = 4
2025-02-20 12:25:04,112 - INFO - Gradient Accumulation steps = 4
2025-02-20 12:25:04,112 - INFO - Total optimization steps = 400

Traceback (most recent call last):
Steps: 2%|▏ | 6/400 [00:27<29:48, 4.54s/it, loss=0.602, lr=1]

File "/workspace/./train_dreambooth_lora_flux.py", line 1935, in
main(args)
File "/workspace/./train_dreambooth_lora_flux.py", line 1837, in main
pipeline = FluxPipeline.from_pretrained(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/dist-packages/huggingface_hub/utils/_validators.py", line 114, in _inner_fn
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/dist-packages/diffusers/pipelines/pipeline_utils.py", line 793, in from_pretrained
config_dict = cls.load_config(cached_folder, dduf_entries=dduf_entries)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/dist-packages/huggingface_hub/utils/_validators.py", line 114, in _inner_fn
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/dist-packages/diffusers/configuration_utils.py", line 381, in load_config
raise EnvironmentError(
OSError: Error no file named model_index.json found in directory /workspace/model/realflux1.

Steps: 2%|▏ | 6/400 [00:28<30:39, 4.67s/it, loss=0.602, lr=1]
Traceback (most recent call last):
File "/usr/local/bin/accelerate", line 8, in
sys.exit(main())
^^^^^^
File "/usr/local/lib/python3.11/dist-packages/accelerate/commands/accelerate_cli.py", line 48, in main
args.func(args)
File "/usr/local/lib/python3.11/dist-packages/accelerate/commands/launch.py", line 1106, in launch_command
simple_launcher(args)
File "/usr/local/lib/python3.11/dist-packages/accelerate/commands/launch.py", line 704, in simple_launcher
raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)

model_index.json

{
  "_class_name": "FluxPipeline",
  "_diffusers_version": "0.32.2",
  "_name_or_path": "black-forest-labs/FLUX.1-schnell",
  "feature_extractor": [
    null,
    null
  ],
  "image_encoder": [
    null,
    null
  ],
  "scheduler": [
    "diffusers",
    "FlowMatchEulerDiscreteScheduler"
  ],
  "text_encoder": [
    "transformers",
    "CLIPTextModel"
  ],
  "text_encoder_2": [
    "transformers",
    "T5EncoderModel"
  ],
  "tokenizer": [
    "transformers",
    "CLIPTokenizer"
  ],
  "tokenizer_2": [
    "transformers",
    "T5TokenizerFast"
  ],
  "transformer": [
    "diffusers",
    "FluxTransformer2DModel"
  ],
  "vae": [
    "diffusers",
    "AutoencoderKL"
  ]
}

[PluginBOXone]

Is it maybe because I am trying to train with the Schnell Version / RealFlux 1.0b_Compact_Schnell? (https://civitai.com/models/788550/realflux-10b)

[sayakpaul]

What is the training command you're using?

[PluginBOXone]

TRAIN_SCRIPT ="./train_dreambooth_lora_flux.py"

command = [
    "accelerate", "launch",
    "--config_file", "/workspace/accelerate.yaml",
    TRAIN_SCRIPT,
    "--pretrained_model_name_or_path", model_name, #MODEL_PATH
    "--instance_data_dir", image_dir,
    "--output_dir", output_dir_final,
    "--instance_prompt", concept_name,
    "--resolution", resolution,
    "--learning_rate", learning_rate,
    "--mixed_precision", mixed_precision
]

of course filled with the right parameters and values from the Dreambooth + LoRA examples

(accelerate launch train_dreambooth_lora_flux.py
--pretrained_model_name_or_path=$MODEL_NAME
--instance_data_dir=$INSTANCE_DIR
--output_dir=$OUTPUT_DIR
--mixed_precision="bf16"
--instance_prompt="a photo of sks dog"
--resolution=512
--train_batch_size=1
--guidance_scale=1
--gradient_accumulation_steps=4
--optimizer="prodigy"
--learning_rate=1.
--report_to="wandb"
--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")

[PluginBOXone]

i really cant explain how this error occours. Since it can clearly start the training, the path and the file structure is right. But somehow it got a problem at Line 1837 and 1887

            pipeline = FluxPipeline.from_pretrained(
                args.pretrained_model_name_or_path,

it cant find the model_index.json on the right path.

As --pretrained_model_name_or_path I pass "/workspace/model/realflux1"

[sayakpaul]

What is the content of /workspace/model/realflux1?

Unfortunately, we cannot reproduce the error like this. Your repository should have the following structure:
https://huggingface.co/black-forest-labs/FLUX.1-dev/tree/main

If you can share the model link on the HF Hub, I can try to see if it's a loading problem.

[PluginBOXone]

I have the model locally on my machine :/ The structure is right, I can inference with no problems with the exact same model/folder. I also noticed, although the training failed after 6 steps as soon as it tried to make a validation image, it still saved a logs folder named "dreambooth-flux-dev-lora" Inside some folders named after i guess random numbers and dots. I found a hparams.yml maybe this can help

adam_beta1: 0.9
adam_beta2: 0.999
adam_epsilon: 1.0e-08
adam_weight_decay: 0.0001
adam_weight_decay_text_encoder: 0.001
allow_tf32: false
cache_dir: null
cache_latents: false
caption_column: null
center_crop: false
checkpointing_steps: 500
checkpoints_total_limit: 1
class_data_dir: null
class_prompt: null
dataloader_num_workers: 0
dataset_config_name: null
dataset_name: null
gradient_accumulation_steps: 4
gradient_checkpointing: false
guidance_scale: 1.0
hub_model_id: null
hub_token: null
image_column: image
instance_data_dir: /workspace/job_files/dada7f5b-d5c6-4ed6-beb4-f84611c78e51-e1/clean_data
instance_prompt: a photo of WIXBSAHA car
learning_rate: 1.0
local_rank: -1
logging_dir: logs
logit_mean: 0.0
logit_std: 1.0
lora_layers: null
lr_num_cycles: 1
lr_power: 1.0
lr_scheduler: constant
lr_warmup_steps: 0
max_grad_norm: 1.0
max_sequence_length: 512
max_train_steps: 400
mixed_precision: bf16
mode_scale: 1.29
num_class_images: 100
num_train_epochs: 67
num_validation_images: 4
optimizer: prodigy
output_dir: /workspace/output_model/dada7f5b-d5c6-4ed6-beb4-e1
pretrained_model_name_or_path: /workspace/model/realflux1
prior_generation_precision: null
prior_loss_weight: 1.0
prodigy_beta3: null
prodigy_decouple: true
prodigy_safeguard_warmup: true
prodigy_use_bias_correction: true
push_to_hub: false
random_flip: false
rank: 4
repeats: 1
report_to: tensorboard
resolution: 768
resume_from_checkpoint: null
revision: null
sample_batch_size: 4
scale_lr: false
seed: 0
text_encoder_lr: 5.0e-06
train_batch_size: 1
train_text_encoder: true
upcast_before_saving: false
use_8bit_adam: false
validation_epochs: 25
validation_prompt: a photo of WIXBSAHA car on a beach
variant: null
weighting_scheme: none
with_prior_preservation: false

is the train script maybe incompatible with the Schnell version? I noticed that the folder name "dreambooth-flux-dev-lora" is with dev, just a guess... The README on HF is also only about the Dev version of FLUX.

[sayakpaul]

Well it seems like you have made some modifications in your setup. So, it is difficult to know for sure what is going wrong.

is the train script maybe incompatible with the Schnell version? I noticed that the folder name "dreambooth-flux-dev-lora" is with dev, just a guess... The README on HF is also only about the Dev version of FLUX.

Technically the script doesn’t handle the timestep distilled nature of Flux Schnell. The error logs you provided suggest something else.

For me to help you further, I need the model to be uploaded on the Hub and I can then try to reproduce it with an example training command from the official README.

[PluginBOXone]

Just made a training test with --pretrained_model_name_or_path "black-forest-labs/FLUX.1-schnell" and now I get this error.

logs (30).txt

[PluginBOXone]

It seems like, you cant train a LoRA directly on the Schnell version...

[sayakpaul]

Can you share the output of:

diffusers-cli env

?