Multiple bugs in flux dreambooth script: train_dreambooth_lora_flux_advanced.py

[freckletonj]

I haven't fully fixed the script, but I'm really not sure how anyone has had success with it (eg the blog post). Many show stoppers when trying to get Textual Inversion working.

To debug, I started peeling away at the main Flux training script, and have fixed all the following bugs, BUT it still doesn't train well. With any hyperparams, it fails to learn a new concept, even to overfit on it, and quickly degrades. So the following are necessary to be fixed, but not sufficient to get this working.

1. On a 1-node-2-gpu setup, accelerator wraps everything in a DistributedSomething class, so all the calls to model.dtype/model.config etc error out, since they're wrapped, and you actually need something like model.module.dtype
2. it requires an --instance-prompt even if you're using a dataset with a custom text column
3. The updated tokenizers are not getting passed into the pipeline
4. text_embeddings get saved weird, it looks like a single copy gets saved no matter the checkpoint number?
5. in log_validation there's a dtype mismatch runtime error when training in fp16/bf16, and you need to autocast
6. also in log_validation the Generator doesn't seem to do anything. Every generation gets performed on a different seed so you can't see the evolution/performance working on the same samples. This is remedied if you save the RNG state of torch/torch_cuda/random, then use the RNG Generator, then reset the random state afterward
7. t5 training doesn't work. On this line it unfreezes token_embedding which doesn't exist in t5. shared needs to be unfrozen.
8. Just a weird one, pulling the embeddings toward std**0.1 makes intuitive sense, but this kind of thing should definitely mention justification. Was this done in a paper somewhere?

Some nice to haves that I wish this did out of the box (but would be happy if it the simple path above "just worked"):

• save latent cache to disk
• aspect ratio bucketing
• 8bit backbone

Honestly, at this point i'm tempted to give up on diffusers and use something else for this current client, after sweating over fixing the above, i still haven't gotten it to work, and still haven't been able to stress test things like the dreambooth functionality, or the whole-text-encoder finetuning feature.

[bonlime]

Half of your links point to SDXL training pipeline, and half to flux training pipeline, is this intended?

[hlky]

cc @linoytsaban @sayakpaul for training

[freckletonj]

@bonlime gah! I was referring between the 2 since the SDXL script does some things right that the flux script doesn't, so accidentally copied the wrong link. I've updated them thanks!

all problems still stand, and I'm happy to help out if someone wants to scour this with me and get it working

[linoytsaban]

Hi @freckletonj, thanks for the detailed issue.

1. that's a good point. @sayakpaul should we make a general refactor?
2. This was done to fall back on when custom captions are missing / column does not exist. That being said, it may indeed be good to make this a non-requirement as briefly discussed here. You can simply pass an empty string in the meantime.
3. For textual inversion inference we only require the trained embeddings to be loaded to the text encoders such that the placeholder tokens correctly map to the learned embeddings. The tokenizers are not optimized.
4. yes. indeed will be better to have this aligned with the transformer checkpoints.
5. which diffusers version are you using? this should fix it -Handling mixed precision for dreambooth flux lora training #9565
6. cc @sayakpaul
7. will look into this

[sayakpaul]

@linoytsaban

that's a good point. @sayakpaul should we make a general refactor?

I think case-by-case is fine. In this case we can use unwrap_model(model).dtype. Should be a simple fix. I would like to also note that in the interest of shipping a script faster, we often just test it on a single GPU and fix it when an error is reported. So, this is fine.

For 6., instead of initializing the generator one time, we can just initialize it everytime pipeline is called when a seed is passed. But regardless, if a generator is initialized only once it will advance with the same states internally and the subsequent generations should still be the same given rest of the things are fixed.

@freckletonj

Some nice to haves that I wish this did out of the box (but would be happy if it the simple path above "just worked"):

For some of these advanced features, you're more than welcome to check out https://github.com/bghira/SimpleTuner and https://github.com/ostris/ai-toolkit/.

We have a flux_lora_quantization that shows how to run DreamBooth + LoRA on a NF4 quantized backbone.

[bonlime]

has anyone even verified that t5 training works? I have tried modifying the script to fix the issues above and tried training for 2000 steps on very small start with TI only. Inspecting the final embeddings and results I get that when loading both embeds for inference, it gets a lot of character likeness. if I load only clip, it gets maybe ~90% there, but with only t5 the results almost don't change compared to no TI. so it makes me wonder has anyone even proved that it's possible to train a "good" t5 embeds with this code?

[hx173149]

Actually, I have found these problem too, I thinkg this example have wasted my two days work.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[gitouni]

I haven't fully fixed the script, but I'm really not sure how anyone has had success with it (eg the blog post). Many show stoppers when trying to get Textual Inversion working.

To debug, I started peeling away at the main Flux training script, and have fixed all the following bugs, BUT it still doesn't train well. With any hyperparams, it fails to learn a new concept, even to overfit on it, and quickly degrades. So the following are necessary to be fixed, but not sufficient to get this working.

1. On a 1-node-2-gpu setup, accelerator wraps everything in a DistributedSomething class, so all the calls to model.dtype/model.config etc error out, since they're wrapped, and you actually need something like model.module.dtype
2. it requires an --instance-prompt even if you're using a dataset with a custom text column
3. The updated tokenizers are not getting passed into the pipeline
4. text_embeddings get saved weird, it looks like a single copy gets saved no matter the checkpoint number?
5. in log_validation there's a dtype mismatch runtime error when training in fp16/bf16, and you need to autocast
6. also in log_validation the Generator doesn't seem to do anything. Every generation gets performed on a different seed so you can't see the evolution/performance working on the same samples. This is remedied if you save the RNG state of torch/torch_cuda/random, then use the RNG Generator, then reset the random state afterward
7. t5 training doesn't work. On this line it unfreezes token_embedding which doesn't exist in t5. shared needs to be unfrozen.
8. Just a weird one, pulling the embeddings toward std**0.1 makes intuitive sense, but this kind of thing should definitely mention justification. Was this done in a paper somewhere?

Some nice to haves that I wish this did out of the box (but would be happy if it the simple path above "just worked"):

• save latent cache to disk
• aspect ratio bucketing
• 8bit backbone

Honestly, at this point i'm tempted to give up on diffusers and use something else for this current client, after sweating over fixing the above, i still haven't gotten it to work, and still haven't been able to stress test things like the dreambooth functionality, or the whole-text-encoder finetuning feature.

Parameters Are Not Updated After optimizer.step

Not only these bugs. I found a fatal one when I tried to check if the parameters in the optimizer are updated after optimizer.step--they're not! That is true even for the embedding parameters of the CLIP (because the t5 parameters are not correctly added into the params_to_optimize)
If you can run the debug launch.json like the following one in the VSCODE, you will find optimizer.param_groups[0]['params'][0] unchanged after optimizer.step().

{
    "version": "0.2.0",
    "configurations": [
        {
            "name": "train_ti_flux_legancy",
            "type": "debugpy",
            "request": "launch",
            "module": "accelerate.commands.launch",
            "console": "integratedTerminal",
            "args": [
                "train_dreambooth_lora_flux_advanced.py",
                "--pretrained_model_name_or_path",
                "pretrained/FLUX.1-dev",
                "--dataset_name",
                "data/img/sea-cucumber",
                "--token_abstraction",
                "sea cucumber",
                "--instance_prompt",
                "sea cucumber",
                "--initializer_concept",
                "seafood",
                "--output_dir",
                "experiments/ip1-Flux-LoRA",
                "--caption_column",
                "clip_prompt",
                "--mixed_precision",
                "bf16",
                "--resolution",
                "512",
                "--train_batch_size",
                "1",
                "--repeats",
                "1",
                "--gradient_accumulation_steps",
                "1",
                "--gradient_checkpointing",
                "--optimizer",
                "AdamW",
                "--train_text_encoder_ti",
                "--enable_t5_ti",
                "--train_text_encoder_ti_frac",
                "1.0",
                "--train_transformer_frac",
                "0",
                "--lr_scheduler",
                "cosine",
                "--lr_warmup_steps",
                "0",
                "--rank",
                "8",
                "--max_train_steps",
                "3000",
                "--checkpointing_steps",
                "500",
                "--seed",
                "0"
            ]
        },
]
}

[rootAvish]

For some of these advanced features, you're more than welcome to check out https://github.com/bghira/SimpleTuner and https://github.com/ostris/ai-toolkit/

The implementation of the custom scheduler in the upstream repo has changed pretty significantly: https://github.com/ostris/ai-toolkit/blob/main/toolkit/samplers/custom_flowmatch_sampler.py with it now containing different schedules from the ones suggested here.

Additionally key functions in the custom scheduler are unused:

• Commented code in the constructor is supposed to consume get_sigmas
• add_noise and get_weights_for_timesteps are presumably for the training loop but never used.

[luchaoqi]

I also got the OP's mentioned bugs and wondering if they gonna be fixed soon.