This repository provides training scripts for Flux model by Black Forest Labs.
XLabs AI team is happy to publish fune-tuning Flux scripts, including:
- LoRA 🔥
- ControlNet 🔥
We trained LoRA and ControlNet models using DeepSpeed!
Both of them are trained on 512x512 pictures, 1024x1024 is in progress.
We trained Canny ControlNet and LoRA checkpoints for FLUX.1 [dev]
You can download them on HuggingFace:
accelerate launch train_flux_lora_deepspeed.py --config "train_configs/test_lora.yaml"accelerate launch train_flux_deepspeed_controlnet.py --config "train_configs/test_canny_controlnet.yaml"Dataset has the following format for the training process:
├── images/
│ ├── 1.png
│ ├── 1.json
│ ├── 2.png
│ ├── 2.json
│ ├── ...
A .json file contains "caption" field with a text prompt.
{
"caption": "A figure stands in a misty landscape, wearing a mask with antlers and dark, embellished attire, exuding mystery and otherworldlines"
}To test our checkpoints, use commands presented below.
prompt: "A girl in a suit covered with bold tattoos and holding a vest pistol, beautiful woman, 25 years old, cool, future fantasy, turquoise & light orange ping curl hair"
prompt: "A handsome man in a suit, 25 years old, cool, futuristic"
python3 main.py \
--prompt "A cute corgi lives in a house made out of sushi, anime" \
--lora_repo_id XLabs-AI/flux-lora-collection --lora_name anime_lora.safetensors \
--device cuda --offload --use_lora --model_type flux-dev-fp8 --width 1024 --height 1024
python3 main.py \
--use_lora --lora_weight 0.7 \
--width 1024 --height 768 \
--lora_repo_id XLabs-AI/flux-lora-collection --lora_name realism_lora.safetensors \
--guidance 4 \
--prompt "contrast play photography of a black female wearing white suit and albino asian geisha female wearing black suit, solid background, avant garde, high fashion"
python3 main.py \
--prompt="a bright blue bird in the garden, natural photo cinematic, MM full HD" \
--repo_id "XLabs-AI/flux-controlnet-canny" \
--name controlnet.safetensors --device cuda --offload --use_controlnet --image "input_image.jpg" --guidance 4
python3 main.py \
--prompt="a dark evil mysterius house with ghosts, cinematic, MM full HD" \
--repo_id "XLabs-AI/flux-controlnet-canny" \
--name controlnet.safetensors --device cuda --offload --use_controlnet --image "input_image.jpg" --guidance 4
python3 main.py \
--prompt="man, 4k photo" \
--repo_id "XLabs-AI/flux-controlnet-canny" \
--name controlnet.safetensors --device cuda --offload --use_controlnet --image "input_image.jpg" --guidance 4
python3 main.py \
--prompt="a oil painting woman sitting at chair and smiling, cinematic, MM full HD" \
--repo_id "XLabs-AI/flux-controlnet-canny" \
--name controlnet.safetensors --device cuda --offload --use_controlnet --image "input_image.jpg" --guidance 4
Use LoRA and Controlnet FP8 version based on Flux-dev-F8 with --offload setting to achieve lower VRAM usage (22 GB) and --name flux-dev-fp8:
python3 main.py \
--offload --name flux-dev-fp8 \
--lora_repo_id XLabs-AI/flux-lora-collection --lora_name realism_lora.safetensors \
--guidance 4 \
--prompt "A handsome girl in a suit covered with bold tattoos and holding a pistol. Animatrix illustration style, fantasy style, natural photo cinematic"
Install our dependencies by running the following command:
pip3 install -r requirements.txtcompute_environment: LOCAL_MACHINE
debug: false
deepspeed_config:
gradient_accumulation_steps: 2
gradient_clipping: 1.0
offload_optimizer_device: none
offload_param_device: none
zero3_init_flag: false
zero_stage: 2
distributed_type: DEEPSPEED
downcast_bf16: 'no'
enable_cpu_affinity: false
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 8
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
Our models fall under the FLUX.1 [dev] Non-Commercial License
Our training and infer scripts under the Apache 2 License
We are working on releasing new ControlNet weight models for Flux: OpenPose, Depth and more!
Stay tuned with XLabs AI to see IP-Adapters for Flux.
