FLUX.1-dev模型本地训练推理教程（GPU/NPU）

模型简介

Flux.1模型是由Black Forest Labs推出的文生图模型套件，具备强大的生成式模型开发能力。该模型由Stable Diffusion的原班人马创立，成员多来自Stable Diffusion的开发团队，因此继承了高水平的模型开发能力。
Flux.1模型包含三个不同的版本：Pro（专业版）、Dev（开发版）和Schnell（本地个人使用版）。

Pro版本提供最先进的性能图像生成，具有顶级的视觉质量、图像细节和输出多样性；
Dev版本是从Pro版本蒸馏而来的，具备相似的图像质量但更高效；
Schnell版本专门为本地开发和个人使用量身定制，生成速度最快，内存占用也最小。

测试环境

GPU：

Centos7.9
CUDA==12.4
torch==2.4.0

NPU:

Centos7.9
CANN==8.0 RC2
torch==2.1.0

推理测试：

NPU/GPU同样

1.克隆代码仓

代码地址：
https://github.com/black-forest-labs/flux#black-forest-labs代码
git clone https://github.com/black-forest-labs/flux
cd /flux
pip install -e ".[all]"

2.下载模型权重

#通过modelscope下载

pip install modelscope
modelscope download --model ai-modelscope/flux.1-dev

3.安装diffusers

pip install git+https://github.com/huggingface/diffusers.git

4.测试脚本

vim test.py
import torch
#from diffusers import FluxPipeline
#####NPU需要添加自动迁移代码
import torch_npu
from torch_npu.contrib import transfer_to_npu
################
from diffusers import FluxPipeline
model_id = "flux.1-dev" #you can also use `black-forest-labs/FLUX.1-dev`
pipe = FluxPipeline.from_pretrained("flux.1-dev",torch_dtype=torch.bfloat16)
pipe = pipe.to('npu')
#pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
prompt = "A cat holding a sign that says hello world"
seed = 42
image = pipe( prompt,
height=512,
width=512, max_sequence_length=128,
output_type="pil",
num_inference_steps=4, #use a larger number if you are using [dev]
generator=torch.Generator("cpu").manual_seed(seed)
).images[0]
image.save("flux-dev.png")

执行脚本进行推理测试

python test.py

在这里插入图片描述

训练测试：

环境准备（使用同一套代码）
1.克隆代码（xlab公布的finetune代码）

代码地址：
https://github.com/XLabs-AI/x-flux/tree/main

git clone https://github.com/XLabs-AI/x-flux.git
cd x-flux/
pip install -r requirements.txt

2.下载模型

#一共需要下载4个模型：flux.1-dev/vae/t5/clip
#使用modelscope下载
#下载flux1-dev和vae
pip install modelscope
modelscope download --model ai-modelscope/flux.1-dev ae.safetensors flux1-dev.safetensors
#下载t5
xlabs-ai/xflux_text_encoders
#下载clip
openai/clip-vit-large-patch14

3.数据集格式

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

{
"caption": "A figure stands in a misty landscape, wearing a mask with antlers and dark, embellished attire, exuding mystery and otherworldlines"
}

在这里插入图片描述

4.config文件

model_name: "../flux-dev"
data_config:
train_batch_size: 1
num_workers: 4
img_size: 512
img_dir: images/
random_ratio: true # support multi crop preprocessing
report_to: wandb
train_batch_size: 1
output_dir: lora/
max_train_steps: 500
learning_rate: 1e-5
lr_scheduler: constant
lr_warmup_steps: 10
adam_beta1: 0.9
adam_beta2: 0.999
adam_weight_decay: 0.01
adam_epsilon: 1e-8
max_grad_norm: 1.0
logging_dir: logs
mixed_precision: "bf16"
checkpointing_steps: 2500
checkpoints_total_limit: 10
tracker_project_name: lora_test
resume_from_checkpoint: latest
gradient_accumulation_steps: 2
rank: 16
single_blocks: "1,2,3,4"
double_blocks: null
disable_sampling: false
sample_every: 1000 # sample every this many steps
sample_width: 1024
sample_height: 1024
sample_steps: 20
sample_prompts:
- "woman with red hair, playing chess at the park, bomb going off in the background"
- "a woman holding a coffee cup, in a beanie, sitting at a cafe"

LoRA微调启动命令：

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

GPU环境训练:
安装依赖：

yum install mesa-libGL.x86_64
pip install -r requirements.txt

拉起任务报错
在这里插入图片描述

config配置文件model name参数，这里不是填写模型权重路径而是类型，因此修改成“flux-dev”

vim ./train_configs/test_lora.yaml

在这里插入图片描述

重新拉起再加载模型文件时报错，
在这里插入图片描述

但相关模型文件已经下载；
查看训练代码发现训练时默认会自动从huggingface上拉取模型文件导致报错，而没有从本地加载模型，查看训练代码需要加载的四个模型，确认本地均已完成下载

vim train_flux_lora_deepspeed.py

在这里插入图片描述

修改模型代码，从本地加载模型，避免从huggingface上拉取，问题解决

vim ./src/flux/util.py

在这里插入图片描述

t5模型和clip模型安装上图所示，创建对应目录，并将模型上传至此即可

拉起任务显卡OOM
在这里插入图片描述

修改img_size值解决

vim train_configs/test_lora.yamlimg_size:512 --->256

在这里插入图片描述

训练到250步时报错；
查看代码每训练250步，会使用wandb深度学习可视化工具用于进行训练结果采集分析，但环境中未安装wandb库，导致import时报错：
在这里插入图片描述

解决wandb方法一：
安装wandb

pip install wandb

拉起任务选择wandb使用模式（3中模式，暂时不清楚怎么使用，选择“3”，本次训练不对训练结果进行处理）
在这里插入图片描述

250步时仍然报wandb初始化配置相关错误;
在这里插入图片描述
尝试修改训练代码去掉wandb工具调用；删除配置文件中wandb超参数字段

删除（注释）训练代码中调用该参数的代码

vim train_flux_lora_deepspeed.py

在这里插入图片描述

删除（注释）训练中采集的代码
在这里插入图片描述
wandb报错问题解决。

解决wandb方法二(推荐)：

修改配置文件，将每250进行一次采集，调整至每1000次（大于训练的总step即可）不影响训练.
在这里插入图片描述

NPU环境训练：

训练代码添加自动迁移代码
在这里插入图片描述

拉起任务

[ERROR] 2024-09-19-15:17:20 (PID:2500, Device:-1, RankID:-1) ERR99999 UNKNOWN application exception
import cv2
File "/usr/local/miniconda3/lib/python3.10/site-packages/cv2/__init__.py", line 181, in <module>
bootstrap()
File "/usr/local/miniconda3/lib/python3.10/site-packages/cv2/__init__.py", line 153, in bootstrap
native_module = importlib.import_module("cv2")
File "/usr/local/miniconda3/lib/python3.10/importlib/__init__.py", line 126, in import_module
return _bootstrap._gcd_import(name[level:], package, level)
ImportError: libGL.so.1: cannot open shared object file: No such file or directory

安装依赖库

yum install mesa-libGL.x86_64
pip install omegaconf
pip install onnxruntime

拉起任务,报缺少optimum-quanto量化的库，
在这里插入图片描述

optimum-quanto库在安装时会自动安装很多依赖，其中包括会自动将torch版本升级至2.4.1版本，使用以下命令仅安装optimum-quanto，防止安装中升级torch及其他依赖库版本

pip install optimum-quanto --no-deps

拉起任务仍然报错，

File "/usr/local/miniconda3/lib/python3.10/site-packages/optimum/quanto/library/extensions/__init__.py", line 17, in <module>
File "/usr/local/miniconda3/lib/python3.10/site-packages/optimum/quanto/__init__.py", line 18, in <module>
File "/usr/local/miniconda3/lib/python3.10/site-packages/optimum/quanto/library/extensions/cpp/__init__.py", line 34, in <module>
from .cpp import *
File "/usr/local/miniconda3/lib/python3.10/site-packages/optimum/quanto/library/extensions/cpp/__init__.py", line 34, in <module>
def unpack_cpp(t: torch.Tensor, bits: int):
File "/usr/local/miniconda3/lib/python3.10/site-packages/torch/library.py", line 155, in wrap
from .library import *
File "/usr/local/miniconda3/lib/python3.10/site-packages/optimum/quanto/library/__init__.py", line 15, in <module>
def unpack_cpp(t: torch.Tensor, bits: int):
File "/usr/local/miniconda3/lib/python3.10/site-packages/torch/library.py", line 155, in wrap
from .extensions import *
File "/usr/local/miniconda3/lib/python3.10/site-packages/optimum/quanto/library/extensions/__init__.py", line 17, in <module>
lib.impl(name, f, dispatch_key)
AttributeError: 'str' object has no attribute 'impl' from .cpp import *
File "/usr/local/miniconda3/lib/python3.10/site-packages/optimum/quanto/library/extensions/cpp/__init__.py", line 34, in <module>
lib.impl(name, f, dispatch_key)
AttributeError
: 'str' object has no attribute 'impl' def unpack_cpp(t: torch.Tensor, bits: int):
File "/usr/local/miniconda3/lib/python3.10/site-packages/torch/library.py", line 155, in wrap
lib.impl(name, f, dispatch_key)
AttributeError: 'str' object has no attribute 'impl'
[ERROR] 2024-09-19-16:00:26 (PID:3032, Device:-1, RankID:-1) ERR99999 UNKNOWN application exception
[ERROR] 2024-09-19-16:00:26 (PID:3028, Device:-1, RankID:-1) ERR99999 UNKNOWN application exception
[ERROR] 2024-09-19-16:00:26 (PID:3029, Device:-1, RankID:-1) ERR99999 UNKNOWN application exception
[ERROR] 2024-09-19-16:00:26 (PID:3026, Device:-1, RankID:-1) ERR99999 UNKNOWN application exception
[ERROR] 2024-09-19-16:00:26 (PID:3031, Device:-1, RankID:-1) ERR99999 UNKNOWN application exception
[ERROR] 2024-09-19-16:00:26 (PID:3030, Device:-1, RankID:-1) ERR99999 UNKNOWN application exception

在GPU环境中正常安装无问题，怀疑是版本问题，由于只安装了optimum-quanto，未安装他需要的依赖库，导致版本有问题，调用时部分功能不支持；

尝试注释调用该库的代码，规避使用optimum-quanto库量化相关库
在这里插入图片描述

在训练时没有调用到，因此先注释掉开头import的地方拉起训练解决，
在这里插入图片描述

拉起训练，显卡OOM
在这里插入图片描述

修改参数

vim train_configs/test_lora.yamlimg_size:256-->128
gradient_accumulation_steps：2-->1
mixed_precision: "bf16" --> mixed_precision: "fp16"

仍然OOM，报错地方调用了pytorch原生适配的FA算子的torch.nn.functional.scaled_dot_product_attention接口
在这里插入图片描述

修改成NPU 优化亲和的FA融合算子torch_npu.npu_fusion_attention

vim ./src/flux/math.py

在这里插入图片描述

拉起训练仍然在NPU FA处报OOM错误同样配置下在GPU环境当中进行训练发现占用显存为76G。因此仅通过优化FA算子无法解决OOM问题
在这里插入图片描述

恢复默认配置，开起ZeRO3功能尝试
配置accelerate库，开启deepspeed ZeRO3特性，
#使用accelerate config命令配置

accelerate config

在这里插入图片描述

配置完成后会在~/.cache/huggingface/accelerate/目录下生成default_config.yaml配置文件（首次配置必须先执行accelerate config生成配置文件，后续可直接修改配置文件即可）

vim ~/.cache/huggingface/accelerate/default_config.yaml

在这里插入图片描述

重新拉起任务问题解决
在这里插入图片描述

模型学习记录