OA0

OA0 是一个探索 AI 的社区

现在注册

已注册用户请登录

OA0 › 代码 › AI Toolkit — 面向模型训练与微调的实用工具集合

AI Toolkit — 面向模型训练与微调的实用工具集合

country · 2026-08-01 11:00:30 · 14 次点击 · 0 条评论

Ostris AI 工具包

AI 工具包是一款易于使用的全能扩散模型训练套件。我致力于在消费级硬件上支持所有最新模型，涵盖图像和视频模型。它既可以作为 GUI 运行，也可以通过 CLI 使用。设计目标是易于使用，同时具备所有可想象的功能。完全免费且开源。

支持的模型

安装

使用 AI 工具包管理器安装（实验性功能）

推荐使用仓库内置的 AI 工具包管理器 来安装和运行 AI 工具包。管理器会检测你的硬件并配置正确的 PyTorch 构建版本，创建 Python 环境，并获取 Node.js 和 FFmpeg 的本地副本——所有内容都保留在 ai-toolkit 文件夹内，不会在系统范围内安装任何东西。每次启动时，管理器都会检查更新并应用它们（你的本地修改绝不会被覆盖——如果你有修改过的文件，更新将被跳过并给出警告），然后启动 UI 于 http://localhost:8675。

管理器目前仍处于实验性阶段——如果遇到任何问题，请告诉我。如果你更喜欢手动安装，下面的手动安装说明仍然有效。

唯一的要求是 git（在 Windows 上，管理器甚至可以获取便携版 git 用于更新，但你需要先安装 git 才能克隆仓库）。

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

然后使用你对应平台的脚本启动管理器：

Linux：

chmod +x run_linux.sh
./run_linux.sh

MacOS（Apple Silicon，实验性）：

chmod +x run_mac.zsh
./run_mac.zsh

Windows：双击 run_windows.bat（或在终端中运行）。

你也可以直接从终端使用管理器（对无头服务器很方便）：

python3 -m manager install   # 首次设置
python3 -m manager update    # 拉取更新 + 同步依赖
python3 -m manager launch    # 启动 UI
python3 -m manager doctor    # 诊断问题

手动安装

要求：
- python >=3.10（推荐 3.12）
- 具有足够显存的 Nvidia GPU 以满足你的需求
- python venv
- git

Linux：

git clone https://github.com/ostris/ai-toolkit.git
cd ai-toolkit
python3 -m venv venv
source venv/bin/activate
# 先安装 torch
pip3 install --no-cache-dir torch==2.13.0 torchvision==0.28.0 torchaudio==2.11.0 --index-url https://download.pytorch.org/whl/cu130
pip3 install -r requirements.txt

对于运行 DGX OS（包括 DGX Spark）的设备，请遵循这些说明。

Windows：

如果你在 Windows 上遇到问题，建议使用 https://github.com/Tavris1/AI-Toolkit-Easy-Install 的简易安装脚本。

git clone https://github.com/ostris/ai-toolkit.git
cd ai-toolkit
python -m venv venv
.\venv\Scripts\activate
pip install --no-cache-dir torch==2.13.0 torchvision==0.28.0 torchaudio==2.11.0 --index-url https://download.pytorch.org/whl/cu130
pip install -r requirements.txt

AI 工具包 UI

AI Toolkit UI

AI 工具包 UI 是 AI 工具包的网页界面。它允许你轻松启动、停止和监控任务。同时支持只需几次点击即可训练模型。你还可以为 UI 设置令牌以防止未经授权的访问，因此在暴露的服务器上运行也比较安全。

运行 UI

要求：
- Node.js > 20

UI 不需要持续运行才能执行训练任务。它仅用于启动/停止/监控任务。以下命令将安装/更新 UI 及其依赖并启动 UI。

cd ui
npm run build_and_start

现在你可以在 http://localhost:8675 访问 UI，如果在服务器上运行，则可通过 http://<你的IP>:8675 访问。

保护 UI 安全

如果你在云服务商或任何不安全的网络上托管 UI，强烈建议使用身份验证令牌进行保护。你可以通过设置环境变量 AI_TOOLKIT_AUTH 来设置超级安全的密码。访问 UI 时需要此令牌。启动 UI 时可以这样设置：

# Linux
AI_TOOLKIT_AUTH=super_secure_password npm run build_and_start

# Windows
set AI_TOOLKIT_AUTH=super_secure_password && npm run build_and_start

# Windows Powershell
$env:AI_TOOLKIT_AUTH="super_secure_password"; npm run build_and_start

训练

将位于 config/examples/train_lora_flux_24gb.yaml 的示例配置文件（schnell 版本使用 config/examples/train_lora_flux_schnell_24gb.yaml）复制到 config 文件夹，并重命名为 你想要的名称.yml
按照文件中的注释编辑文件
运行文件：python run.py config/你想要的名称.yml

启动时会在配置文件中指定的位置创建一个以训练名称命名的文件夹。其中包含所有检查点和图像。您可以随时使用 ctrl+c 停止训练，恢复时会从最后一个检查点继续。

重要提示：如果在保存时按下 ctrl+c，很可能会损坏该检查点。请等到保存完成后再停止。

需要帮助？

除非是代码中的错误，否则请不要开 bug 报告。欢迎加入我的 Discord 并在那里寻求帮助。但请避免直接 PM 我提出一般性问题或寻求支持。请在 Discord 中提问，我会在有空时回答。

Ostris Cloud

您可以使用许多云服务商来租用 GPU。如果想以最大的方式支持这个项目，请考虑使用 Ostris Cloud。Ostris Cloud 由我（Ostris）拥有和运营，每一美元的收入都直接用于资助这个项目的开发。

在 RunPod 中训练

如果您想使用 Runpod，但还没有注册，请考虑使用我的 Runpod 联盟链接来支持这个项目。

我维护了一个官方的 Runpod Pod 模板，可以在此访问。

我还制作了一个简短视频，展示如何开始使用 AI 工具包搭配 Runpod，点击观看。

在 Modal 中训练

1. 设置

ai-toolkit：

git clone https://github.com/ostris/ai-toolkit.git
cd ai-toolkit
git submodule update --init --recursive
python -m venv venv
source venv/bin/activate
pip install torch
pip install -r requirements.txt
pip install --upgrade accelerate transformers diffusers huggingface_hub #可选，遇到问题时运行

Modal：

运行 pip install modal 安装 modal Python 包。
运行 modal setup 进行身份验证（如果不起作用，请尝试 python -m modal setup）。

Hugging Face：

从这里获取 READ 令牌，并从这里申请 Flux.1-dev 模型的访问权限。
运行 huggingface-cli login 并粘贴你的令牌。

2. 上传数据集

将包含 .jpg、.jpeg 或 .png 图像和 .txt 文件的数据集文件夹拖放到 ai-toolkit 中。

3. 配置

将位于 config/examples/modal 的示例配置文件复制到 config 文件夹，并重命名为 你想要的名称.yml。
按照文件中的注释编辑配置，请务必小心并遵循示例中的 /root/ai-toolkit 路径。

4. 编辑 run_modal.py

在 code_mount = modal.Mount.from_local_dir 中设置你本地完整的 ai-toolkit 路径，例如：

code_mount = modal.Mount.from_local_dir("/Users/username/ai-toolkit", remote_path="/root/ai-toolkit")
- 在 @app.function 中选择 GPU 和 Timeout（默认为 A100 40GB 和 2 小时超时）。

5. 训练

在你的终端中运行配置文件：modal run run_modal.py --config-file-list-str=/root/ai-toolkit/config/你想要的名称.yml。
你可以在本地终端或 modal.com 上监控训练进度。
模型、样本和优化器将存储在 Storage > flux-lora-models 中。

6. 保存模型

通过运行 modal volume ls flux-lora-models 查看卷的内容。
通过运行 modal volume get flux-lora-models 你的模型名称 下载内容。
示例：modal volume get flux-lora-models my_first_flux_lora_v1。

Modal 截图

Modal 训练截图

数据集准备

数据集通常需要是一个包含图像和相应文本文件的文件夹。目前仅支持 jpg、jpeg 和 png 格式。Webp 格式目前存在问题。文本文件应与图像同名，但扩展名为 .txt。例如 image2.jpg 和 image2.txt。文本文件应只包含说明文字（caption）。你可以在说明文件中添加 [trigger] 这个词，如果配置中有 trigger_word，它会被自动替换。

图像不会被放大，但会被缩小并放入批次桶中。你不需要裁剪/调整图像大小。加载器会自动调整它们的大小，并能处理不同的宽高比。

训练特定层

要使用 LoRA 训练特定层，可以使用 only_if_contains 网络参数。例如，如果你想仅训练 The Last Ben 使用的 2 层（在这篇帖子中提到），你可以这样调整网络参数：

      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"

层的命名约定采用 diffusers 格式，因此查看模型的状态字典可以找出要训练的层的名称后缀。你也可以使用此方法仅训练特定的权重组。例如，仅训练 FLUX.1 的 single_transformer，可以使用以下配置：

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

你也可以通过 ignore_if_contains 网络参数排除某些层。例如，要排除所有 single transformer 块：

      network:
        type: "lora"
        linear: 128
        linear_alpha: 128
        network_kwargs:
          ignore_if_contains:
            - "transformer.single_transformer_blocks."

ignore_if_contains 的优先级高于 only_if_contains。因此，如果一个权重同时被两者覆盖，它将被忽略。

LoKr 训练

要了解更多关于 LoKr 的信息，请在 KohakuBlueleaf/LyCORIS 阅读更多内容。要训练 LoKr 模型，可以像这样调整配置文件中的网络类型：

      network:
        type: "lokr"
        lokr_full_rank: true
        lokr_factor: 8

其他一切（包括层定位）都应正常工作。

支持我的工作

如果你喜欢我的项目或将其用于商业用途，请考虑赞助我。每一份帮助都很重要！💖

当前赞助商

以下所有个人/组织都是无私地使这个项目成为可能的人。谢谢你们！！

项目地址：https://github.com/ostris/ai-toolkit

14 次点击 ∙ 0 人收藏

登录后收藏

0 条回复