Ostris AI 工具包

AI 工具包是一款易于使用的、一体化的扩散模型训练套件。我致力于在消费级硬件上支持所有最新的模型，涵盖图像和视频领域。该工具包既可以通过图形用户界面运行，也可以通过命令行界面操作。它设计得简单易用，同时具备你能想到的所有功能。完全免费且开源。

支持的模型

图像

• black-forest-labs/FLUX.1-dev (FLUX.1)
• black-forest-labs/FLUX.2-dev (FLUX.2)
• black-forest-labs/FLUX.2-klein-base-4B (FLUX.2-klein-base-4B)
• black-forest-labs/FLUX.2-klein-base-9B (FLUX.2-klein-base-9B)
• ostris/Flex.1-alpha (Flex.1)
• ostris/Flex.2-preview (Flex.2)
• lodestones/Chroma1-Base (Chroma)
• Alpha-VLLM/Lumina-Image-2.0 (Lumina2)
• Qwen/Qwen-Image (Qwen-Image)
• Qwen/Qwen-Image-2512 (Qwen-Image-2512)
• HiDream-ai/HiDream-I1-Full (HiDream)
• OmniGen2/OmniGen2 (OmniGen2)
• Tongyi-MAI/Z-Image-Turbo (Z-Image Turbo)
• Tongyi-MAI/Z-Image (Z-Image)
• ostris/Z-Image-De-Turbo (Z-Image De-Turbo)
• stabilityai/stable-diffusion-xl-base-1.0 (SDXL)
• stable-diffusion-v1-5/stable-diffusion-v1-5 (SD 1.5)

指令 / 编辑

• black-forest-labs/FLUX.1-Kontext-dev (FLUX.1-Kontext-dev)
• Qwen/Qwen-Image-Edit (Qwen-Image-Edit)
• Qwen/Qwen-Image-Edit-2509 (Qwen-Image-Edit-2509)
• Qwen/Qwen-Image-Edit-2511 (Qwen-Image-Edit-2511)
• HiDream-ai/HiDream-E1-1 (HiDream E1)

视频

• Wan-AI/Wan2.1-T2V-1.3B-Diffusers (Wan 2.1 1.3B)
• Wan-AI/Wan2.1-I2V-14B-480P-Diffusers (Wan 2.1 I2V 14B-480P)
• Wan-AI/Wan2.1-I2V-14B-720P-Diffusers (Wan 2.1 I2V 14B-720P)
• Wan-AI/Wan2.1-T2V-14B-Diffusers (Wan 2.1 14B)
• Wan-AI/Wan2.2-T2V-A14B-Diffusers (Wan 2.2 14B)
• Wan-AI/Wan2.2-I2V-A14B-Diffusers (Wan 2.2 I2V 14B)
• Wan-AI/Wan2.2-TI2V-5B-Diffusers (Wan 2.2 TI2V 5B)
• Lightricks/LTX-2 (LTX-2)
• Lightricks/LTX-2.3 (LTX-2.3)

实验性

• lodestones/Zeta-Chroma (Zeta Chroma)

安装

要求：

• Python >=3.10（推荐使用 3.12）
• 具备足够显存的 NVIDIA GPU
• Python 虚拟环境
• Git

Linux:

git clone https://github.com/ostris/ai-toolkit.git
cd ai-toolkit
python3 -m venv venv
source venv/bin/activate
# 首先安装 PyTorch
pip3 install --no-cache-dir torch==2.9.1 torchvision==0.24.1 torchaudio==2.9.1 --index-url https://download.pytorch.org/whl/cu128
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.9.1 torchvision==0.24.1 torchaudio==2.9.1 --index-url https://download.pytorch.org/whl/cu128
pip install -r requirements.txt

MacOS:

目前提供对 Apple Silicon Mac 的实验性支持。由于我没有配备足够内存的 Mac 来进行全面测试，因此如果你遇到任何问题，请及时告知我。我们提供了一个便捷的脚本来在 macOS 上安装并运行工具包，路径为 ./run_mac.zsh，该脚本会本地安装依赖项并启动 UI。要运行此脚本，请执行以下步骤：

git clone https://github.com/ostris/ai-toolkit.git
cd ai-toolkit
chmod +x run_mac.zsh
./run_mac.zsh

AI 工具包 UI

AI 工具包 UI 是一款用于 AI 工具包的网页界面。它允许你轻松地启动、停止和监控任务，只需点击几下即可训练模型。此外，你还可以为 UI 设置一个访问令牌，以防止未经授权的访问，因此在暴露的服务器上运行也相对安全。

运行 UI

要求：

• Node.js > 20

UI 并不需要一直运行才能让任务执行。它仅用于启动、停止和监控任务。以下命令将安装或更新 UI 及其依赖项，并启动 UI。

cd ui
npm run build_and_start

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

保护 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

训练

1. 将位于 config/examples/train_lora_flux_24gb.yaml（schnell 版本为 config/examples/train_lora_flux_schnell_24gb.yaml）的示例配置文件复制到 config 文件夹，并重命名为 whatever_you_want.yml。
2. 按照文件中的注释编辑该文件。
3. 运行该文件：python run.py config/whatever_you_want.yml。

启动训练时，会创建一个与配置文件中指定名称相同的文件夹，其中包含所有检查点和图像。您可以随时按 Ctrl+C 停止训练，下次继续时将从最后一个检查点恢复。

重要提示：如果在保存检查点时按下 Ctrl+C，很可能会导致该检查点损坏。因此，请务必等待保存完成后再停止。

需要帮助吗？

请仅在代码存在 bug 时才提交 bug 报告。欢迎加入我的 Discord 服务器 https://discord.gg/VXmU2f5WEU 寻求帮助。不过，请不要直接私信我询问一般性问题或寻求支持。请在 Discord 中提问，我会在方便时回复。

Gradio 界面

要在本地使用自定义界面开始训练，按照上述步骤操作并安装好 ai-toolkit 后：

cd ai-toolkit # 如果您尚未进入 ai-toolkit 文件夹
huggingface-cli login # 提供一个具有写入权限的 token，以便在训练结束后发布您的 LoRA
python flux_train_ui.py

您将启动一个 UI，允许您上传图片、添加描述、训练并发布您的 LoRA。

在 RunPod 上训练

如果您想使用 RunPod 但尚未注册，请考虑通过我的 RunPod 联盟链接 https://runpod.io?ref=h0y9jyr2 来支持该项目。

我在此维护了一个官方的 RunPod Pod 模板，可通过 这里 访问。

我还制作了一段简短的视频，介绍如何使用 AI Toolkit 和 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：

• 从 这里 获取一个读取权限的 token，并从 这里 申请访问 Flux.1-dev 模型的权限。
• 运行 huggingface-cli login 并粘贴您的 token。

2. 上传数据集

• 将包含 .jpg、.jpeg 或 .png 图片以及 .txt 文件的数据集文件夹拖放到 ai-toolkit 目录中。

3. 配置

• 将位于 config/examples/modal 的示例配置文件复制到 config 文件夹，并重命名为 whatever_you_want.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/whatever_you_want.yml。
• 您可以在本地终端或 modal.com 上监控训练进度。
• 模型、样本和优化器将存储在 Storage > flux-lora-models 中。

6. 保存模型

• 运行 modal volume ls flux-lora-models 查看卷的内容。
• 运行 modal volume get flux-lora-models your-model-name 下载内容。
• 示例：modal volume get flux-lora-models my_first_flux_lora_v1。

Modal 截图

---

数据集准备

数据集通常需要是一个包含图片及其对应文本文件的文件夹。目前仅支持 jpg、jpeg 和 png 格式，webp 格式暂不支持。文本文件应与图片同名，但扩展名为 .txt。例如，image2.jpg 和 image2.txt。文本文件中只需包含描述文字。您可以在描述文件中添加 [trigger] 关键字，如果您的配置中设置了 trigger_word，它将会被自动替换。

图片不会被放大，但会被缩小并分组以进行批量处理。您无需裁剪或调整图片大小。加载器会自动调整图片尺寸，并能处理不同宽高比的图片。

训练特定层

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

      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_blocks，

      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

其他设置应保持不变，包括层的选择等。

支持我的工作

如果你喜欢我的项目，或者在商业项目中使用了它们，请考虑资助我。每一份支持都对我很有帮助！💖

当前赞助者

所有这些个人和组织都在无私地支持这个项目的开展。非常感谢你们！！

最后更新：2026年3月31日 18:10 UTC

---

---

---

---

---

Ostris AI Toolkit 快速上手指南

Ostris AI Toolkit 是一款专为消费级硬件设计的全能扩散模型训练套件，支持图像和视频模型的训练。它提供图形界面（GUI）和命令行（CLI）两种操作模式，旨在让复杂的模型训练变得简单易用。

1. 环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux, Windows, 或 macOS (Apple Silicon 实验性支持)
• Python: 版本 >= 3.10 (推荐 3.12)
• GPU: NVIDIA GPU (需具备足够的显存以运行目标模型)
• 其他依赖:
  • git
  • Python 虚拟环境工具 (venv)
  • UI 模式额外需要: Node.js > 20

2. 安装步骤

Linux / Windows 通用安装

1. 克隆仓库并进入目录

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

2. 创建并激活虚拟环境

   • Linux:

     python3 -m venv venv
     source venv/bin/activate
     

   • Windows (CMD):

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

   • Windows (PowerShell):

     python -m venv venv
     .\venv\Scripts\Activate.ps1
     

3. 安装 PyTorch (CUDA 12.8) 注意：此处使用官方源以确保兼容性，国内用户若遇网络问题可配置 pip 镜像源（如清华源），但 PyTorch 二进制包建议优先保证版本一致。

   pip install --no-cache-dir torch==2.9.1 torchvision==0.24.1 torchaudio==2.9.1 --index-url https://download.pytorch.org/whl/cu128
   

4. 安装项目依赖

   pip install -r requirements.txt
   

macOS (Apple Silicon) 安装

macOS 用户可使用提供的便捷脚本自动安装依赖并启动 UI：

git clone https://github.com/ostris/ai-toolkit.git
cd ai-toolkit
chmod +x run_mac.zsh
./run_mac.zsh

安装 Web UI (可选)

如果您希望使用网页界面来监控和管理训练任务：

cd ui
npm run build_and_start

启动后访问 http://localhost:8675。若在服务器运行，可将 localhost 替换为服务器 IP。

3. 基本使用

方式一：命令行训练 (CLI)

这是最核心的使用方式，通过配置文件启动训练。

1. 准备配置文件 将示例配置文件复制到 config 文件夹并重命名（例如 my_train.yml）：

   # 示例：复制 FLUX.1 的 LoRA 训练配置
   cp config/examples/train_lora_flux_24gb.yaml config/my_train.yml
   

2. 编辑配置 打开 config/my_train.yml，根据注释修改参数。主要需关注：

   • model: 选择支持的模型路径（如 black-forest-labs/FLUX.1-dev）。
   • dataset: 指向包含图片和对应 .txt 标签文件的文件夹路径。
   • output: 设置模型保存路径。
   • 提示：图片无需预先裁剪或调整大小，加载器会自动处理分桶和缩放。

3. 开始训练 运行以下命令启动训练：

   python run.py config/my_train.yml
   

   • 训练过程中可随时按 Ctrl+C 停止，程序会自动保存当前检查点并在下次运行时从中断处恢复。
   • 注意：请勿在正在写入检查点时强制终止，以免文件损坏。

方式二：Gradio 简易界面

适合本地快速微调并发布 LoRA 模型。

1. 登录 Hugging Face

   huggingface-cli login
   

   (输入具有 write 权限的 Token)

2. 启动 UI

   python flux_train_ui.py
   

   在浏览器中打开显示的地址，即可上传图片、编写提示词、训练并直接发布 LoRA 模型。

数据集格式说明

无论使用哪种方式，数据集需遵循以下结构：

• 文件夹内包含图片文件 (.jpg, .jpeg, .png)。
• 每张图片需有一个同名的 .txt 文件作为标签（例如 image1.jpg 对应 image1.txt）。
• .txt 文件内容仅为该图片的描述文本（Caption）。
• 可在文本中使用 [trigger] 占位符，并在配置文件中定义 trigger_word 进行自动替换。