用于微调扩散模型的终极训练工具包
AI Toolkit is an easy to use all in one training suite for diffusion models. I try to support all the latest models on consumer grade hardware. Image and video models. It can be run as a GUI or CLI. It is designed to be easy to use but still have every feature imaginable. Free and open source.
The recommended way to install and run AI Toolkit is with the AI Toolkit
Manager, built into this repo. The manager detects your hardware and sets up
the right PyTorch build, creates the python environment, and grabs local copies
of Node.js and FFmpeg — everything stays inside the ai-toolkit folder, nothing
is installed system-wide. On every launch the manager checks for updates and
applies them (your local changes are never overwritten — if you have modified
files, the update is skipped with a warning), then starts the UI at
http://localhost:8675.
The manager is still experimental — please let me know if you have any issues with it. The manual instructions below still work if you prefer them or run into problems.
The only requirement is git (on Windows the manager can even fetch a portable git for updates, but you need one installed to clone the repo first).
git clone https://github.com/ostris/ai-toolkit.git
cd ai-toolkit
Then start the manager with the script for your platform:
Linux:
chmod +x run_linux.sh
./run_linux.sh
MacOS (Apple Silicon, experimental):
chmod +x run_mac.zsh
./run_mac.zsh
Windows: double-click run_windows.bat (or run it from a terminal).
You can also use the manager directly from a terminal (handy on headless servers):
python3 -m manager install # first-time setup
python3 -m manager update # pull updates + sync dependencies
python3 -m manager launch # start the UI
python3 -m manager doctor # diagnose problems
Requirements:
Linux:
git clone https://github.com/ostris/ai-toolkit.git
cd ai-toolkit
python3 -m venv venv
source venv/bin/activate
# install torch first
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
For devices running DGX OS (including DGX Spark), follow these instructions.
Windows:
If you are having issues with Windows. I recommend using the easy install script at 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
The AI Toolkit UI is a web interface for the AI Toolkit. It allows you to easily start, stop, and monitor jobs. It also allows you to easily train models with a few clicks. It also allows you to set a token for the UI to prevent unauthorized access so it is mostly safe to run on an exposed server.
Requirements:
The UI does not need to be kept running for the jobs to run. It is only needed to start/stop/monitor jobs. The commands below will install / update the UI and it's dependencies and start the UI.
cd ui
npm run build_and_start
You can now access the UI at http://localhost:8675 or http://<your-ip>:8675 if you are running it on a server.
If you are hosting the UI on a cloud provider or any network that is not secure, I highly recommend securing it with an auth token.
You can do this by setting the environment variable AI_TOOLKIT_AUTH to super secure password. This token will be required to access
the UI. You can set this when starting the UI like so:
# 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 (config/examples/train_lora_flux_schnell_24gb.yaml for schnell) to the config folder and rename it to whatever_you_want.ymlpython run.py config/whatever_you_want.ymlA folder with the name and the training folder from the config file will be created when you start. It will have all checkpoints and images in it. You can stop the training at any time using ctrl+c and when you resume, it will pick back up from the last checkpoint.
IMPORTANT. If you press crtl+c while it is saving, it will likely corrupt that checkpoint. So wait until it is done saving
Please do not open a bug report unless it is a bug in the code. You are welcome to Join my Discord and ask for help there. However, please refrain from PMing me directly with general question or support. Ask in the discord and I will answer when I can.
You can use many cloud providers to rent GPUs. If you want to help support this project in the largest way possible, please consider using Ostris Cloud. Ostris Cloud is owned and operated by me, Ostris, and every dollar earned goes directly back into funding the development of this project.
If you would like to use Runpod, but have not signed up yet, please consider using my Runpod affiliate link to help support this project.
I maintain an official Runpod Pod template here which can be accessed here.
I have also created a short video showing how to get started using AI Toolkit with Runpod here.
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 #Optional, run it if you run into issues
pip install modal to install the modal Python package.modal setup to authenticate (if this doesn’t work, try `python -m modal暂无开放 Issues，或尚未同步最近议题。