AI Toolkit：各种 AI 脚本，主要用于 Stable Diffusion

ostris/ai-toolkit 是一个 GitHub 仓库， 包含多种 AI 脚本，主要用于 Stable Diffusion 相关任务。该项目使用 MIT 许可证，目前已获得 4.2k star, 470 forks, 59 watching。

主要功能：

  • Stable Diffusion 脚本: 提供各种用于 Stable Diffusion 的脚本和工具。
  • FLUX.1 模型训练: 支持对 FLUX.1 系列模型进行训练，包括 FLUX.1-dev 和 FLUX.1-schnell。
  • LoRA 和 LoKr 训练: 支持 LoRA (Low-Rank Adaptation) 和 LoKr 训练，允许针对特定层进行训练。
  • Gradio UI: 提供一个基于 Gradio 的用户界面，方便用户上传图片、标注、训练和发布 LoRA 模型。
  • UI界面: 基于Node.js 和 React 搭建的WEB UI界面， 它可以让你方便地启动、停止、监控训练任务，也能更容易地开始新的训练模型。该UI还处于早期beta阶段，可能存在bug 并且UI安全性未知，不建议暴露到公网。

安装步骤：

通用准备：

  • Python 3.10+
  • Nvidia GPU (所需的显存大小取决于具体任务)
  • python venv
  • git

Linux:

git clone https://github.com/ostris/ai-toolkit.git
cd ai-toolkit
git submodule update --init --recursive
python3 -m venv venv
source venv/bin/activate
# .\venv\Scripts\activate on windows
# install torch first
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
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 >= 18

你每次更新代码的时候都需要按照如下步骤重新执行构建：

cd ui
npm install
npm run build
npm run update_db

运行UI界面

确保你已经完成了构建步骤，UI界面不需要保持运行才能使任务正常执行，它只用于启动、停止和监控任务。

cd ui
npm run start

现在你可以通过http://localhost:8675 或者 http://你的IP:8675 访问UI。

FLUX.1 训练:

  • 硬件需求: 至少 24GB 显存的 GPU。
  • FLUX.1-dev: 采用非商业许可，需要先在 Hugging Face 上接受模型访问许可，并配置 HF_TOKEN 环境变量。
  • FLUX.1-schnell: 采用 Apache 2.0 许可，不需要 HF_TOKEN，但需要特殊的训练适配器。

训练步骤:

  1. 复制示例配置文件 (config/examples/train_lora_flux_24gb.yaml 或 config/examples/train_lora_flux_schnell_24gb.yaml) 到 config 文件夹并重命名。
  2. 按照文件中的注释编辑配置文件。
  3. 运行命令 python run.py config/your_config_file.yml。

RunPod 和 Modal 训练:

文档还提供了在 RunPod 和 Modal 上进行训练的详细步骤， 包括环境配置、数据集上传、配置文件的修改以及训练命令。

数据集准备:

数据集需要包含图片和对应的文本文件，文本文件与图片文件同名，扩展名为 .txt，内容为图片的 caption。 支持的图片格式包括 jpg, jpeg, 和 png。

训练特定层:

可以通过配置 network_kwargs 中的 only_if_contains 或 ignore_if_contains 来指定训练的层。

LoKr 训练:

可以通过修改配置文件中的 network 类型为 "lokr" 来进行 LoKr 训练。

总结:

ai-toolkit 提供了一套完整的 AI 脚本，尤其是在 Stable Diffusion 模型训练和微调方面， 提供了多种配置选项和训练方法。它也支持使用UI界面进行训练。

希望对大家有所帮助！