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)
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 ..
- FLUX denoiser and FLUX vae from here
- clip_l.safetensors and T5_fp16.safetensors from here
Easiest way to download these models is:
pip install huggingface_hub
- Grab your Huggingface token (account --> settings --> Access Tokens)
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
- Create a folder of training images
- make a copy of
template/train_config.json
and adjust with your training setup. - Optionally adjust
template/eval_prompts.txt
- run
python main.py --config /path/to/train_config.json
- All the logs, samples and .safetensors files will appear under ./results
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