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 🔥

Training

We trained LoRA and ControlNet models using DeepSpeed!
Both of them are trained on 512x512 pictures, 1024x1024 is in progress.

Models

We trained Canny ControlNet and LoRA checkpoints for FLUX.1 [dev]
You can download them on HuggingFace:

• flux-controlnet-canny
• flux-RealismLora
• flux-lora-collections

LoRA

accelerate launch train_flux_lora_deepspeed.py --config "train_configs/test_lora.yaml"

ControlNet

accelerate launch train_flux_deepspeed_controlnet.py --config "train_configs/test_canny_controlnet.yaml"

Training Dataset

Dataset has the following format for the training process:

├── images/
│    ├── 1.png
│    ├── 1.json
│    ├── 2.png
│    ├── 2.json
│    ├── ...

Example images/*.json file

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"
}

Inference

To test our checkpoints, use commands presented below.

LoRA

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"

ControlNet (Canny)

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

Low memory mode

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"

Requirements

Install our dependencies by running the following command:

pip3 install -r requirements.txt

Accelerate Configuration Example

compute_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

Models Licence

Our models fall under the FLUX.1 [dev] Non-Commercial License
Our training and infer scripts under the Apache 2 License

Near Updates

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.