FLUX-trainer

Flux LoRA trainer and full-finetuning.

For now this is a utility wrapper around kohya that deals with:

• dataset preparation and cleaning
• automatic captioning (using Florence2)
• easily passing in training args through config.json files
• running sample inference using sample prompts provided in a .txt file
• packaging and uploading outputs into a .tar file to upload (TODO)

Setup instructions:

1. Setup the environment:

conda create --name flux python=3.10
conda activate flux
git clone https://github.com/edenartlab/flux-trainer.git
cd flux-trainer
pip install -r requirements.txt
git clone https://github.com/kohya-ss/sd-scripts.git
cd sd-scripts
git checkout a2ad7e5644f08141fe053a2b63446d70d777bdcf
git checkout sd3
pip install -r requirements.txt
cd ..

2. Download the Flux models into ./models:

• FLUX denoiser and FLUX vae from here
• clip_l.safetensors and T5_fp16.safetensors from here

Easiest way to download these models is:

1. pip install huggingface_hub
2. Grab your Huggingface token (account --> settings --> Access Tokens)
3. huggingface-cli login

And then run:

mkdir models
cd models
huggingface-cli download black-forest-labs/FLUX.1-dev ae.safetensors --repo-type model --local-dir .
huggingface-cli download black-forest-labs/FLUX.1-dev flux1-dev.safetensors  --repo-type model --local-dir .
huggingface-cli download comfyanonymous/flux_text_encoders clip_l.safetensors --repo-type model --local-dir .
huggingface-cli download comfyanonymous/flux_text_encoders t5xxl_fp16.safetensors --repo-type model --local-dir .

If you have all these models already downloaded somewhere else, you can just point to their paths in your train_config.json

3. Run a training job:

1. Create a folder of training images
2. make a copy of template/train_config.json and adjust with your training setup.
3. Optionally adjust template/eval_prompts.txt
4. run python main.py --config /path/to/train_config.json
5. All the logs, samples and .safetensors files will appear under ./results

Building with docker:

docker build --build-arg HF_TOKEN=your_hf_token -t flux-trainer .

or

docker build --no-cache --build-arg HF_TOKEN=your_hf_token -t flux-trainer .

and then run eg: docker run -it flux-trainer

or get a shell inside the container: docker run -it flux-trainer /bin/bash