364 lines
17 KiB
Markdown
364 lines
17 KiB
Markdown
# Ostris AI Toolkit
|
||
|
||
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.
|
||
|
||
|
||
|
||
## Supported Models
|
||
|
||
### Image
|
||
- [black-forest-labs/FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) (FLUX.1)
|
||
- [black-forest-labs/FLUX.2-dev](https://huggingface.co/black-forest-labs/FLUX.2-dev) (FLUX.2)
|
||
- [black-forest-labs/FLUX.2-klein-base-4B](https://huggingface.co/black-forest-labs/FLUX.2-klein-base-4B) (FLUX.2-klein-base-4B)
|
||
- [black-forest-labs/FLUX.2-klein-base-9B](https://huggingface.co/black-forest-labs/FLUX.2-klein-base-9B) (FLUX.2-klein-base-9B)
|
||
- [ostris/Flex.1-alpha](https://huggingface.co/ostris/Flex.1-alpha) (Flex.1)
|
||
- [ostris/Flex.2-preview](https://huggingface.co/ostris/Flex.2-preview) (Flex.2)
|
||
- [lodestones/Chroma1-Base](https://huggingface.co/lodestones/Chroma1-Base) (Chroma)
|
||
- [Alpha-VLLM/Lumina-Image-2.0](https://huggingface.co/Alpha-VLLM/Lumina-Image-2.0) (Lumina2)
|
||
- [Qwen/Qwen-Image](https://huggingface.co/Qwen/Qwen-Image) (Qwen-Image)
|
||
- [Qwen/Qwen-Image-2512](https://huggingface.co/Qwen/Qwen-Image-2512) (Qwen-Image-2512)
|
||
- [HiDream-ai/HiDream-I1-Full](https://huggingface.co/HiDream-ai/HiDream-I1-Full) (HiDream I1)
|
||
- [OmniGen2/OmniGen2](https://huggingface.co/OmniGen2/OmniGen2) (OmniGen2)
|
||
- [Tongyi-MAI/Z-Image-Turbo](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo) (Z-Image Turbo)
|
||
- [Tongyi-MAI/Z-Image](https://huggingface.co/Tongyi-MAI/Z-Image) (Z-Image)
|
||
- [ostris/Z-Image-De-Turbo](https://huggingface.co/ostris/Z-Image-De-Turbo) (Z-Image De-Turbo)
|
||
- [zhen-nan/L2P](https://huggingface.co/zhen-nan/L2P) (Z-Image L2P)
|
||
- [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) (SDXL)
|
||
- [stable-diffusion-v1-5/stable-diffusion-v1-5](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) (SD 1.5)
|
||
- [baidu/ERNIE-Image](https://huggingface.co/baidu/ERNIE-Image) (ERNIE-Image)
|
||
- [NucleusAI/Nucleus-Image](https://huggingface.co/NucleusAI/Nucleus-Image) (Nucleus-Image)
|
||
- [Boogu/Boogu-Image-0.1-Base](https://huggingface.co/Boogu/Boogu-Image-0.1-Base) (Boogu Image 0.1)
|
||
- [HiDream-ai/HiDream-O1-Image](https://huggingface.co/HiDream-ai/HiDream-O1-Image) (HiDream O1)
|
||
- [ideogram-ai/ideogram-4-fp8](https://huggingface.co/ideogram-ai/ideogram-4-fp8) (Ideogram 4 FP8)
|
||
- [Photoroom/prxpixel-t2i](https://huggingface.co/Photoroom/prxpixel-t2i) (PRXPixel)
|
||
- [circlestone-labs/Anima-Base-v1.0-Diffusers](https://huggingface.co/circlestone-labs/Anima-Base-v1.0-Diffusers) (Anima)
|
||
- [krea/Krea-2-Raw](https://huggingface.co/krea/Krea-2-Raw) (Krea 2)
|
||
- [krea/Krea-2-Turbo](https://huggingface.co/krea/Krea-2-Turbo) (Krea 2 Turbo)
|
||
- [microsoft/Mage-Flow-Base](https://huggingface.co/microsoft/Mage-Flow-Base) (Mage-Flow)
|
||
|
||
### Instruction / Edit
|
||
- [black-forest-labs/FLUX.1-Kontext-dev](https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev) (FLUX.1-Kontext-dev)
|
||
- [Qwen/Qwen-Image-Edit](https://huggingface.co/Qwen/Qwen-Image-Edit) (Qwen-Image-Edit)
|
||
- [Qwen/Qwen-Image-Edit-2509](https://huggingface.co/Qwen/Qwen-Image-Edit-2509) (Qwen-Image-Edit-2509)
|
||
- [Qwen/Qwen-Image-Edit-2511](https://huggingface.co/Qwen/Qwen-Image-Edit-2511) (Qwen-Image-Edit-2511)
|
||
- [HiDream-ai/HiDream-E1-1](https://huggingface.co/HiDream-ai/HiDream-E1-1) (HiDream E1)
|
||
- [Boogu/Boogu-Image-0.1-Edit](https://huggingface.co/Boogu/Boogu-Image-0.1-Edit) (Boogu Image Edit)
|
||
- [krea/Krea-2-Raw](https://huggingface.co/krea/Krea-2-Raw) (Krea 2 Edit Training)
|
||
- [krea/Krea-2-Turbo](https://huggingface.co/krea/Krea-2-Turbo) (Krea 2 Turbo Edit Training)
|
||
- [microsoft/Mage-Flow-Edit-Base](https://huggingface.co/microsoft/Mage-Flow-Edit-Base) (Mage-Flow Edit)
|
||
|
||
### Video
|
||
- [Wan-AI/Wan2.1-T2V-1.3B-Diffusers](https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B-Diffusers) (Wan 2.1 1.3B)
|
||
- [Wan-AI/Wan2.1-I2V-14B-480P-Diffusers](https://huggingface.co/Wan-AI/Wan2.1-I2V-14B-480P-Diffusers) (Wan 2.1 I2V 14B-480P)
|
||
- [Wan-AI/Wan2.1-I2V-14B-720P-Diffusers](https://huggingface.co/Wan-AI/Wan2.1-I2V-14B-720P-Diffusers) (Wan 2.1 I2V 14B-720P)
|
||
- [Wan-AI/Wan2.1-T2V-14B-Diffusers](https://huggingface.co/Wan-AI/Wan2.1-T2V-14B-Diffusers) (Wan 2.1 14B)
|
||
- [Wan-AI/Wan2.2-T2V-A14B-Diffusers](https://huggingface.co/Wan-AI/Wan2.2-T2V-A14B-Diffusers) (Wan 2.2 14B)
|
||
- [Wan-AI/Wan2.2-I2V-A14B-Diffusers](https://huggingface.co/Wan-AI/Wan2.2-I2V-A14B-Diffusers) (Wan 2.2 I2V 14B)
|
||
- [Wan-AI/Wan2.2-TI2V-5B-Diffusers](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B-Diffusers) (Wan 2.2 TI2V 5B)
|
||
- [Lightricks/LTX-2](https://huggingface.co/Lightricks/LTX-2) (LTX-2)
|
||
- [Lightricks/LTX-2.3](https://huggingface.co/Lightricks/LTX-2.3) (LTX-2.3)
|
||
- [MiniMaxAI/MiniMax-H3](https://huggingface.co/MiniMaxAI/MiniMax-H3) (MiniMaxAI/MiniMax-H3)
|
||
|
||
### Audio
|
||
- [ACE-Step/Ace-Step1.5](https://huggingface.co/ACE-Step/Ace-Step1.5) (Ace Step 1.5)
|
||
- [ACE-Step/acestep-v15-xl-base](https://huggingface.co/ACE-Step/acestep-v15-xl-base) (Ace Step 1.5 XL)
|
||
|
||
### Experimental
|
||
- [lodestones/Zeta-Chroma](https://huggingface.co/lodestones/Zeta-Chroma) (Zeta Chroma)
|
||
|
||
## Installation
|
||
|
||
### Install with the AI Toolkit Manager (experimental)
|
||
|
||
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).
|
||
|
||
```bash
|
||
git clone https://github.com/ostris/ai-toolkit.git
|
||
cd ai-toolkit
|
||
```
|
||
|
||
Then start the manager with the script for your platform:
|
||
|
||
Linux:
|
||
```bash
|
||
chmod +x run_linux.sh
|
||
./run_linux.sh
|
||
```
|
||
|
||
MacOS (Apple Silicon, experimental):
|
||
```bash
|
||
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):
|
||
|
||
```bash
|
||
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
|
||
```
|
||
|
||
### Manual installation
|
||
|
||
Requirements:
|
||
- python >=3.10 (3.12 recommended)
|
||
- Nvidia GPU with enough ram to do what you need
|
||
- python venv
|
||
- git
|
||
|
||
|
||
Linux:
|
||
```bash
|
||
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](dgx_instructions.md) 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](https://github.com/Tavris1/AI-Toolkit-Easy-Install)
|
||
|
||
```bash
|
||
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 Toolkit UI
|
||
|
||
<img src="https://ostris.com/wp-content/uploads/2025/02/toolkit-ui.jpg" alt="AI Toolkit UI" width="100%">
|
||
|
||
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.
|
||
|
||
## Running the UI
|
||
|
||
Requirements:
|
||
- Node.js > 20
|
||
|
||
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.
|
||
|
||
```bash
|
||
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.
|
||
|
||
## Securing the UI
|
||
|
||
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:
|
||
|
||
```bash
|
||
# 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
|
||
```
|
||
|
||
### Training
|
||
1. Copy the example config file located at `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.yml`
|
||
2. Edit the file following the comments in the file
|
||
3. Run the file like so `python run.py config/whatever_you_want.yml`
|
||
|
||
A 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
|
||
|
||
### Need help?
|
||
|
||
Please do not open a bug report unless it is a bug in the code. You are welcome to [Join my Discord](https://discord.gg/VXmU2f5WEU)
|
||
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.
|
||
|
||
## Ostris Cloud
|
||
|
||
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](https://cloud.ostris.com). Ostris Cloud is owned and operated by me, Ostris, and every dollar earned goes directly back into funding the development of this project.
|
||
|
||
<a href="https://cloud.ostris.com" target="_blank"><img src="https://cloud.ostris.com/api/og" alt="Ostris Cloud" style="max-width:100%;width:600px;height:auto;"></a>
|
||
|
||
|
||
## Training in RunPod
|
||
If you would like to use Runpod, but have not signed up yet, please consider using [my Runpod affiliate link](https://runpod.io?ref=h0y9jyr2) to help support this project.
|
||
|
||
|
||
I maintain an official Runpod Pod template here which can be accessed [here](https://console.runpod.io/deploy?template=0fqzfjy6f3&ref=h0y9jyr2).
|
||
|
||
I have also created a short video showing how to get started using AI Toolkit with Runpod [here](https://youtu.be/HBNeS-F6Zz8).
|
||
|
||
## Training in Modal
|
||
|
||
### 1. Setup
|
||
#### 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 #Optional, run it if you run into issues
|
||
```
|
||
#### Modal:
|
||
- Run `pip install modal` to install the modal Python package.
|
||
- Run `modal setup` to authenticate (if this doesn’t work, try `python -m modal setup`).
|
||
|
||
#### Hugging Face:
|
||
- Get a READ token from [here](https://huggingface.co/settings/tokens) and request access to Flux.1-dev model from [here](https://huggingface.co/black-forest-labs/FLUX.1-dev).
|
||
- Run `huggingface-cli login` and paste your token.
|
||
|
||
### 2. Upload your dataset
|
||
- Drag and drop your dataset folder containing the .jpg, .jpeg, or .png images and .txt files in `ai-toolkit`.
|
||
|
||
### 3. Configs
|
||
- Copy an example config file located at ```config/examples/modal``` to the `config` folder and rename it to ```whatever_you_want.yml```.
|
||
- Edit the config following the comments in the file, **<ins>be careful and follow the example `/root/ai-toolkit` paths</ins>**.
|
||
|
||
### 4. Edit run_modal.py
|
||
- Set your entire local `ai-toolkit` path at `code_mount = modal.Mount.from_local_dir` like:
|
||
|
||
```
|
||
code_mount = modal.Mount.from_local_dir("/Users/username/ai-toolkit", remote_path="/root/ai-toolkit")
|
||
```
|
||
- Choose a `GPU` and `Timeout` in `@app.function` _(default is A100 40GB and 2 hour timeout)_.
|
||
|
||
### 5. Training
|
||
- Run the config file in your terminal: `modal run run_modal.py --config-file-list-str=/root/ai-toolkit/config/whatever_you_want.yml`.
|
||
- You can monitor your training in your local terminal, or on [modal.com](https://modal.com/).
|
||
- Models, samples and optimizer will be stored in `Storage > flux-lora-models`.
|
||
|
||
### 6. Saving the model
|
||
- Check contents of the volume by running `modal volume ls flux-lora-models`.
|
||
- Download the content by running `modal volume get flux-lora-models your-model-name`.
|
||
- Example: `modal volume get flux-lora-models my_first_flux_lora_v1`.
|
||
|
||
### Screenshot from Modal
|
||
|
||
<img width="1728" alt="Modal Traning Screenshot" src="https://github.com/user-attachments/assets/7497eb38-0090-49d6-8ad9-9c8ea7b5388b">
|
||
|
||
---
|
||
|
||
## Dataset Preparation
|
||
|
||
Datasets generally need to be a folder containing images and associated text files. Currently, the only supported
|
||
formats are jpg, jpeg, and png. Webp currently has issues. The text files should be named the same as the images
|
||
but with a `.txt` extension. For example `image2.jpg` and `image2.txt`. The text file should contain only the caption.
|
||
You can add the word `[trigger]` in the caption file and if you have `trigger_word` in your config, it will be automatically
|
||
replaced.
|
||
|
||
Images are never upscaled but they are downscaled and placed in buckets for batching. **You do not need to crop/resize your images**.
|
||
The loader will automatically resize them and can handle varying aspect ratios.
|
||
|
||
|
||
## Training Specific Layers
|
||
|
||
To train specific layers with LoRA, you can use the `only_if_contains` network kwargs. For instance, if you want to train only the 2 layers
|
||
used by The Last Ben, [mentioned in this post](https://x.com/__TheBen/status/1829554120270987740), you can adjust your
|
||
network kwargs like so:
|
||
|
||
```yaml
|
||
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"
|
||
```
|
||
|
||
The naming conventions of the layers are in diffusers format, so checking the state dict of a model will reveal
|
||
the suffix of the name of the layers you want to train. You can also use this method to only train specific groups of weights.
|
||
For instance to only train the `single_transformer` for FLUX.1, you can use the following:
|
||
|
||
```yaml
|
||
network:
|
||
type: "lora"
|
||
linear: 128
|
||
linear_alpha: 128
|
||
network_kwargs:
|
||
only_if_contains:
|
||
- "transformer.single_transformer_blocks."
|
||
```
|
||
|
||
You can also exclude layers by their names by using `ignore_if_contains` network kwarg. So to exclude all the single transformer blocks,
|
||
|
||
|
||
```yaml
|
||
network:
|
||
type: "lora"
|
||
linear: 128
|
||
linear_alpha: 128
|
||
network_kwargs:
|
||
ignore_if_contains:
|
||
- "transformer.single_transformer_blocks."
|
||
```
|
||
|
||
`ignore_if_contains` takes priority over `only_if_contains`. So if a weight is covered by both,
|
||
if will be ignored.
|
||
|
||
## LoKr Training
|
||
|
||
To learn more about LoKr, read more about it at [KohakuBlueleaf/LyCORIS](https://github.com/KohakuBlueleaf/LyCORIS/blob/main/docs/Guidelines.md). To train a LoKr model, you can adjust the network type in the config file like so:
|
||
|
||
```yaml
|
||
network:
|
||
type: "lokr"
|
||
lokr_full_rank: true
|
||
lokr_factor: 8
|
||
```
|
||
|
||
Everything else should work the same including layer targeting.
|
||
|
||
|
||
## Support My Work
|
||
|
||
If you enjoy my projects or use them commercially, please consider sponsoring me. Every bit helps! 💖
|
||
|
||
<a href="https://ostris.com/sponsors" target="_blank"><img src="https://ostris.com/wp-content/uploads/2025/05/support-banner2.png" alt="Support my work" style="max-width:100%;height:auto;"></a>
|
||
|
||
### Current Sponsors
|
||
|
||
All of these people / organizations are the ones who selflessly make this project possible. Thank you!!
|
||
|
||
<a href="https://ostris.com/sponsors"><img src="https://ostris.com/sponsors.svg" alt="Sponsors" style="width:100%;height:auto;"></a>
|