训练

我们使用DeepSpeed训练了LoRA和ControlNet模型！ 两者都在512x512像素的图片上进行了训练，1024x1024的训练正在进行中。

模型

我们为FLUX.1 [dev]训练了Canny ControlNet和LoRA检查点 你可以在HuggingFace上下载它们：

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"

训练数据集

数据集的训练过程格式如下：

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

`images/*.json`文件示例

.json文件包含带有文本提示的"caption"字段。

{
    "caption": "一个人物站在雾蒙蒙的风景中，戴着带有鹿角的面具，穿着黑色的华丽服装，散发出神秘和超凡脱俗的气息"
}

推理

要测试我们的检查点，请使用下面展示的命令。

LoRA

示例图片1 提示词："一个穿着西装的女孩，身上布满大胆的纹身，手持背心手枪，美丽的女性，25岁，酷炫，未来幻想，蓝绿色和浅橙色的卷发" 示例图片2 提示词："一个穿着西装的英俊男人，25岁，酷炫，未来风格"

python3 main.py \
 --prompt "带有'hello world'文字的女性兽人精灵" \
 --lora_repo_id XLabs-AI/flux-furry-lora --lora_name furry_lora.safetensors --device cuda --offload --use_lora \
 --model_type flux-dev-fp8 --width 1024 --height 1024 \
 --timestep_to_start_cfg 1 --num_steps 25 --true_gs 3.5 --guidance 4

示例图片1

python3 main.py \
--prompt "一只可爱的柯基住在由寿司做成的房子里，动漫风格" \
--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

示例图片3

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 "对比摄影，一位穿着白色西装的黑人女性和一位穿着黑色西装的白化病亚洲艺伎女性，纯色背景，前卫，高级时尚"

示例图片3

ControlNet (Canny)

python3 main.py \
    --prompt="花园里一只明亮的蓝色鸟，自然照片电影风格，MM全高清" \
    --repo_id "XLabs-AI/flux-controlnet-canny" \
    --name controlnet.safetensors --device cuda --offload --use_controlnet --image "input_image.jpg" --guidance 4 --control_type canny

示例图片1

python3 main.py \
    --prompt="一座黑暗邪恶神秘的鬼屋，电影风格，MM全高清" \
    --repo_id "XLabs-AI/flux-controlnet-canny" \
    --name controlnet.safetensors --device cuda --offload --use_controlnet --image "input_image.jpg" --guidance 4 --control_type canny

示例图片2

python3 main.py \
    --prompt="男人，4K照片" \
    --repo_id "XLabs-AI/flux-controlnet-canny" \
    --name controlnet.safetensors --device cuda --offload --use_controlnet --image "input_image.jpg" --guidance 4 --control_type canny

示例图片3

python3 main.py \
    --prompt="一幅油画，一位女士坐在椅子上微笑，电影风格，MM全高清" \
    --repo_id "XLabs-AI/flux-controlnet-canny" \
    --name controlnet.safetensors --device cuda --offload --use_controlnet --image "input_image.jpg" --guidance 4 --control_type canny

示例图片4

低内存模式

使用基于Flux-dev-F8的LoRA和Controlnet FP8版本，设置--offload以实现更低的显存使用（22 GB），并使用--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 "一个身材姣好的女孩穿着西装，身上布满大胆的纹身，手持手枪。动画矩阵插图风格，奇幻风格，自然照片电影风格"

示例图片0

依赖要求

通过运行以下命令安装我们的依赖：

pip3 install -r requirements.txt

Accelerate配置示例

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