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

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 demo_lora_inference.py \
    --repo_id XLabs-AI/flux-RealismLora \
    --width 1024 --height 768
    --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"

Low memory mode

Use LoRA FP8 version based on Flux-dev-F8 with --offload setting to achieve lower VRAM usage (22 GB):

python3 demo_lora_inference.py \
    --repo_id XLabs-AI/flux-RealismLora \
    --prompt "A handsome girl in a suit covered with bold tattoos and holding a pistol. Animatrix illustration style, fantasy style, natural photo cinematic" --offload --name flux-dev-fp8

ControlNet (Canny)

python3 demo_controlnet_inference.py \
    --checkpoint controlnet.safetensors \
    --control_image "input_image.jpg" \
    --prompt "a bright blue bird in the garden, natural photo cinematic, MM full HD"

python3 demo_controlnet_inference.py \
    --checkpoint controlnet.safetensors \
    --control_image "input_image.jpg" \
    --prompt "a dark evil mysterius house with ghosts, cinematic, MM full HD"

python3 demo_controlnet_inference.py \
    --checkpoint controlnet.safetensors \
    --control_image "input_image.jpg" \
    --prompt "a handsome viking man with white hair, cinematic, MM full HD"

python3 demo_controlnet_inference.py \
    --checkpoint controlnet.safetensors \
    --control_image "input_image.jpg" \
    --prompt "a oil painting woman sitting at chair and smiling, cinematic, MM full HD"

Requirements

Install our dependencies by running the following command:

pip3 install 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.