AI Toolkit by Ostris：Stable Diffusion与FLUX.1模型训练工具包-首席AI分享圈

🚀邀请体验：中国首家 AI IDE 智能编程软件 Trae 中文版下载，不限量使用 DeepSeek-R1 和 Doubao-pro!

综合介绍

AI Toolkit by Ostris 是一个开源的AI工具集，专注于支持Stable Diffusion及FLUX.1模型的训练与图像生成任务。该工具集由开发者Ostris创建并维护，托管于GitHub，旨在为研究者和开发者提供灵活的模型微调与实验平台。它包含多种AI脚本，支持LoRA提取、图像批量生成及特定层训练等功能。目前项目处于开发阶段，部分功能可能不够稳定，但其高度可定制性使其适用于深度学习领域的高级用户。工具集支持Linux和Windows系统，运行FLUX.1模型训练需至少24GB显存的Nvidia GPU。

AI Toolkit by Ostris：Stable Diffusion与FLUX.1模型训练工具包-1

功能列表

模型训练: 支持Stable Diffusion及FLUX.1模型微调，可训练LoRA和LoKr模型。
图像生成: 根据配置文件或文本提示生成批量图像。
LoRA提取与优化: 提供LoRA和LoCON提取工具，优化模型特征提取。
特定层训练: 可选择特定神经网络层进行训练，灵活调整权重。
用户界面支持: 提供AI Toolkit UI和Gradio UI，简化任务管理和模型训练操作。
数据集处理: 自动调整图像分辨率并按桶分组，支持多种图像格式。
云端训练: 支持在RunPod和Modal平台运行训练任务。

使用帮助

安装流程

Linux系统安装

克隆仓库: 在终端运行以下命令下载代码：

git clone https://github.com/ostris/ai-toolkit.git
cd ai-toolkit

更新子模块: 确保所有依赖库完整：

git submodule update --init --recursive

创建虚拟环境: 使用Python 3.10或更高版本：

python3 -m venv venv
source venv/bin/activate

安装依赖: 先安装PyTorch，再安装其他依赖：

pip3 install torch
pip3 install -r requirements.txt

Windows系统安装

克隆仓库: 在命令提示符中运行：

git clone https://github.com/ostris/ai-toolkit.git
cd ai-toolkit

更新子模块:

git submodule update --init --recursive

创建虚拟环境:

python -m venv venv
.\venv\Scripts\activate

安装依赖: 安装支持CUDA 12.4的PyTorch版本，再安装其他依赖：

pip install torch==2.5.1 torchvision==0.20.1 --index-url https://download.pytorch.org/whl/cu124
pip install -r requirements.txt

UI界面安装

安装Node.js: 确保系统安装Node.js 18或更高版本。
构建UI: 进入ui目录并安装依赖：

cd ui
npm install
npm run build
npm run update_db

运行UI: 启动界面：

npm run start

访问UI: 在浏览器中输入http://localhost:8675。

主要功能操作流程

FLUX.1模型训练

准备环境: 确保GPU显存至少24GB。若用于显示输出，需在配置文件中设置low_vram: true，以在CPU上量化模型。
配置FLUX.1-dev:

登录Hugging Face，访问black-forest-labs/FLUX.1-dev并接受许可。
在项目根目录创建.env文件，添加HF_TOKEN=你的读取密钥。

配置FLUX.1-schnell:

编辑配置文件（如train_lora_flux_schnell_24gb.yaml），添加：

model:
name_or_path: "black-forest-labs/FLUX.1-schnell"
assistant_lora_path: "ostris/FLUX.1-schnell-training-adapter"
is_flux: true
quantize: true
sample:
guidance_scale: 1
sample_steps: 4

准备数据集: 在根目录创建dataset文件夹，放入.jpg、.jpeg或.png图像及对应的.txt描述文件。
编辑配置文件: 复制config/examples/train_lora_flux_24gb.yaml到config目录，重命名为my_config.yml，修改folder_path为数据集路径。
运行训练: 执行：

python run.py config/my_config.yml

训练结果保存于指定输出文件夹，可用Ctrl+C暂停并从最近检查点恢复。

使用Gradio UI训练

登录Hugging Face: 运行huggingface-cli login，输入具有write权限的密钥。
启动UI: 执行：

python flux_train_ui.py

操作UI: 在界面中上传图像、填写描述、设置参数后点击训练，完成后可发布LoRA模型。

云端训练（RunPod）

创建RunPod实例: 使用模板runpod/pytorch:2.2.0-py3.10-cuda12.1.1-devel-ubuntu22.04，选择A40（48GB显存）。
安装工具集: 连接Jupyter Notebook，在终端运行Linux安装命令。
上传数据集: 在根目录创建dataset文件夹，拖入图像和描述文件。
配置并运行: 修改配置文件中的folder_path，执行python run.py config/my_config.yml。

数据集准备

格式要求: 支持.jpg、.jpeg、.png格式，描述文件为.txt，文件名需一致（如image1.jpg对应image1.txt）。
描述内容: .txt中写入描述，可使用[trigger]占位符，由配置文件中的trigger_word替换。
自动调整: 工具根据配置分辨率自动缩小并分组图像，不支持放大。

特定层训练

编辑配置文件: 在network部分添加：

network:
type: "lora"
linear: 128
linear_alpha: 128
network_kwargs:
only_if_contains:
- "transformer.single_transformer_blocks.7.proj_out"
- "transformer.single_transformer_blocks.20.proj_out"