AI-Toolkit LoRA Training Guide: FLUX.2, Z-Image & Qwen-Image on Your Own GPU

If you want to train a LoRA for FLUX.2, Z-Image or Qwen-Image, use ostris/ai-toolkit — the MIT-licensed trainer (11.6K GitHub stars as of August 2026) that covers the entire new model wave, while kohya's sd-scripts stops at FLUX.1. Its reference configs put FLUX.1-dev and even the 20B Qwen-Image on a single 24GB card — rank 16, lr 1e-4, batch 1, 2,000 steps, quantized base — and nothing in its examples folder targets less than 24GB.

That paragraph is the guide in miniature. The rest is the detail, verified against the repository itself in August 2026: the real VRAM table built from the repo's own reference configs (not vibes), the install commands that currently work, the dataset rules the loader actually enforces, and the limits — including the models that have code support but no published config yet. One thing you will not find here is training-time estimates. The repo publishes none, they swing wildly with card, dataset, resolution and steps, and we don't print numbers we can't stand behind.

---

Why AI-Toolkit Won the Trainer Race {#why-ai-toolkit}

AI-Toolkit is where support for new image and video models lands first — and for the 2026 wave, it is effectively the only open-source place it landed at all.

Check the competition's own paperwork. The supported-architecture list in kohya's sd-scripts README covers SD 1.x/2.x, SDXL, SD3/3.5, FLUX.1, Lumina, HunyuanImage-2.1 and Anima. No FLUX.2. No Z-Image. No Qwen-Image. No Wan. If your target model shipped after mid-2025, the kohya route mostly ends at a GitHub issue asking when support is coming.

AI-Toolkit's README, meanwhile, lists trainable support spanning FLUX.1 (dev, schnell, Kontext), FLUX.2 (dev plus the klein-base 4B and 9B), Chroma, Lumina2, Qwen-Image and Qwen-Image-Edit, HiDream, OmniGen2, Z-Image (Turbo, base, De-Turbo and L2P variants), SDXL and SD 1.5 on the image side — plus video models (Wan 2.1/2.2 from 1.3B to 14B, LTX-2, MiniMax-H3) and even audio. One tool, one dataset format, one config schema across all of them.

A few more reasons it became the default, all verifiable from the repo:

• It is genuinely free. MIT license, active development (pushed as recently as early August 2026). The author funds it through an optional cloud service and a RunPod template — the local tool is not a demo for a paid product.
• It has a real web UI. A local dashboard at http://localhost:8675 for building jobs, watching sample images and managing datasets — and the UI does not need to stay running for jobs to keep training.
• Checkpointing is sane. Stop any run with ctrl+c and it resumes from the last checkpoint (the README's one warning: don't hit ctrl+c while it is saving).
• Cloud escape hatches are first-party. An official RunPod template and Modal configs ship in the repo, so the exact same config file runs on rented hardware when your card is too small.

If you are still on the SDXL/FLUX.1 stack, our image LoRA training guide covers that era's tooling (kohya included) and goes deeper on dataset curation — the two guides are complementary, and the dataset advice there applies here unchanged.

---

VRAM by Model: The Real Table {#vram-requirements}

A 24GB card covers almost every reference config the project ships. The two exceptions go up, not down: Qwen-Image-Edit targets 32GB and HiDream 48GB.

This table is not our estimate. It is read directly off the example configs in the repo's config/examples/ folder — ostris names the files after the VRAM tier they were built for, which makes them the most honest public statement of what actually fits.

Source: ostris/ai-toolkit config/examples/ directory, retrieved August 2026. Parameter counts from the respective official model cards.

Two important absences: FLUX.2 and Z-Image have full training support in the codebase (each has its own module under extensions_built_in/diffusion_models/ — flux2/ and z_image/ — alongside qwen_image/, wan22/, ltx2/ and the rest) but no published reference config yet. For those two you build the job in the web UI instead of copying a yaml. What we can honestly infer about their tiers:

• Z-Image is the light one. It is a 6B model — half the size of FLUX.1's 12B, which trains comfortably in 24GB — and its AI-Toolkit implementation exposes the same quantization, low_vram and layer-offloading knobs as the heavier models. If any model of this wave ends up trainable on a 16GB card, it is this one; Tongyi's own model card already puts inference "comfortably within 16G VRAM consumer devices." We have not verified a 16GB training run, so treat that as the experiment to try, not a spec.
• FLUX.2-dev is the heavy one. It is a 32B rectified-flow transformer per the Black Forest Labs model card. When the 20B Qwen-Image already needs 3-bit quantization to fit 24GB, a 32B model will not need less. Treat 24GB as the absolute floor with aggressive quantization and offloading, and 32GB+ (RTX 5090 tier) or rented hardware as the comfortable path.

If you are choosing hardware for this workload, our best GPU for image generation ranking covers the buy decision, and the RTX 3090 local AI guide makes the case for the cheapest used 24GB entry ticket.

---

Install and First Run {#install}

Fastest path: clone the repo and run the launcher script for your OS — it sets up everything and serves the web UI at http://localhost:8675.

git clone https://github.com/ostris/ai-toolkit.git
cd ai-toolkit
./run_linux.sh     # macOS: ./run_mac.zsh · Windows: double-click run_windows.bat

Prefer to manage the environment yourself? The manual Linux install from the README (Python 3.10 minimum, 3.12 recommended):

python3 -m venv venv
source venv/bin/activate
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 the web UI on its own (Node.js > 20 required — the same command installs, updates and starts it):

cd ui
npm run build_and_start

Then open http://localhost:8675. If the machine is reachable from outside your LAN, set the AI_TOOLKIT_AUTH environment variable — it becomes the token the UI asks for. Jobs keep running with the UI closed.

Headless/CLI route — the workflow the README describes: copy an example config out of config/examples/, rename it, edit it following the comments in the file, then:

python run.py config/whatever_you_want.yml

Ctrl+c stops the run; starting it again resumes from the last checkpoint. The one way to hurt yourself: interrupting during a checkpoint save can corrupt that checkpoint — watch the console before you kill it.

---

Dataset Prep the Way the Loader Wants It {#dataset}

One folder of jpg/jpeg/png images, each with a same-named .txt caption file. Do not crop, do not resize — the loader buckets and downscales automatically, and it never upscales.

The rules, straight from the docs:

• Formats: only jpg, jpeg and png are currently supported. Convert the webp screenshots first.
• Captions: image2.jpg pairs with image2.txt. The text file is the caption, nothing else.
• Trigger word: write [trigger] inside captions and set trigger_word in your config — it is substituted automatically at train time. This is how you make a LoRA answer to "photo of [trigger] at the beach" without hand-editing every caption.
• Aspect ratios: mixed sizes and ratios are fine. Images are sorted into resolution buckets for batching; oversized images are downscaled, small ones are left alone (never upscaled — so genuinely tiny images just waste a slot).

The repo deliberately does not prescribe a dataset size, and dataset curation is where most LoRAs are actually won or lost — quality and variety beat count. Our image LoRA training guide covers curation and captioning strategy in detail; all of it transfers directly.

---

The Settings That Matter {#settings}

Every ostris reference config starts from the same baseline: LoRA rank 16 / alpha 16, lr 1e-4, adamw8bit, batch size 1, 2,000 steps, resolution buckets [512, 768, 1024], quantized base model, flowmatch scheduler. Change things one at a time from there.

Reading the FLUX.1-dev 24GB config tells you what a first run looks like:

• steps: 2000 and lr: 1e-4 — the stock recipe. Sample images render every 250 steps and a checkpoint saves every 250 steps, so you get eight looks at the LoRA as it cooks and can stop at whichever checkpoint nailed it. That is the honest way to pick step count: watch the samples, don't trust a formula.
• linear: 16, linear_alpha: 16 — rank 16 is the default capacity. Character/subject LoRAs usually live happily here; raise it when a complex style refuses to stick, at the cost of VRAM and file size.
• quantize: true — the base model is quantized during training on the 24GB tier. This is normal and expected; it is the reason a 12B model plus gradients fits a consumer card at all.
• gradient_checkpointing: true, batch_size: 1, optimizer: adamw8bit — the standard memory-saving trio. Leave them alone on 24GB.
• Save format is float16 safetensors — a standard LoRA file that drops straight into ComfyUI or any loader for the matching base.

The Qwen-Image 24GB config shows how far the project pushes when the model is 20B: the transformer is quantized to 3-bit (torchao uint3) with an ostris-published accuracy recovery adapter layered on to claw quality back, the text encoder runs at qfloat8, low_vram: true is on, and text-embedding caching is enabled with the comment "required for 24GB." You don't need to understand every trick — but you should understand that the yaml is the VRAM budget. Raise the resolution, the rank, or the batch size and you have left the envelope the filename promised.

---

FLUX.2, Z-Image and Qwen-Image: Model Notes {#model-notes}

Short version: Z-Image is the cheap experiment, Qwen-Image is the proven 24GB workhorse, FLUX.2 is the quality play that wants the most VRAM — and the licenses differ more than the VRAM does.

Z-Image (6B, Apache 2.0)

Tongyi's single-stream DiT (per its model card) is the smallest serious training target of the wave, and the AI-Toolkit README lists four trainable variants: Turbo (Tongyi-MAI/Z-Image-Turbo), the undistilled base (Tongyi-MAI/Z-Image — "a full-capacity, undistilled transformer," in Tongyi's words), De-Turbo, and L2P. Worth knowing before you train: Turbo is an 8-step distilled model, and distilled models are historically awkward LoRA targets — the distillation objective fights the fine-tune. The tell is that De-Turbo is published under ostris' own Hugging Face account (ostris/Z-Image-De-Turbo): the trainer's author shipped a de-distilled Z-Image precisely because people want to train on it. If your Turbo-trained LoRA degrades the model's speed advantage or picks up artifacts, the undistilled variants are the fix. Already generating with it? Our Z-Image Turbo ComfyUI guide covers the inference side, including the 16GB/8GB/6GB tiers.

Qwen-Image (20B, Apache 2.0)

The best-documented big-model path in the repo: a maintained 24GB reference config (the 3-bit story above), plus a separate Qwen-Image-Edit config at the 32GB tier for training edit-instruction LoRAs. Apache licensing means what you train is yours to use commercially — a real consideration if LoRA work is client work. Background on the model itself in our Qwen-Image local guide.

FLUX.2 (32B, FLUX Non-Commercial License on dev weights)

The strongest image quality of the three and the heaviest lift. Training support is in the codebase, configuration happens in the web UI, and there is no published VRAM reference yet — our reasoning in the table section applies: 24GB floor, 32GB+ comfortable. The README also lists the smaller FLUX.2-klein-base-4B and klein-base-9B as trainable, and the 4B klein is the interesting one for modest cards: per its model card it is Apache 2.0 and fits in ~13GB of VRAM for inference — no training reference exists yet, so treat trainability on small cards as unproven. Note the license asymmetry on the flagship: FLUX.2-dev weights ship under the FLUX Non-Commercial License per the model card, so check terms before commercial LoRA work — or pick Chroma, which is Apache 2.0 per its model card and has its own 24GB reference config. Inference-side setup lives in our FLUX.2 local guide.

---

Video LoRAs: Wan 2.2 and Friends {#video}

The Wan 2.2 14B reference config targets the same 24GB tier as the image models — video LoRA training on a consumer card is no longer exotic.

That train_lora_wan22_14b_24gb.yaml file is quietly the most remarkable thing in the examples folder: a 14B video model, trainable on the same used RTX 3090 that trains your FLUX LoRAs. A Wan 2.1 1.3B config sits alongside it for faster iteration, and the README's supported list extends to LTX-2 and MiniMax-H3 (UI-configured, like the newest image models). Same dataset format, same config schema — a video LoRA is not a new skill, just a longer wait.

---

No 24GB Card? Your Honest Options {#low-vram}

Below 24GB: experiment locally with Z-Image, or rent exactly the hardware the tool expects — the repo ships an official RunPod template, and its Modal setup defaults to an A100 40GB.

• 16GB cards: Z-Image (6B) is the one new-wave model plausibly within reach, for the size reasons above — unverified, but cheap to try. Everything else in the table is out of local reach at reference settings.
• 8-12GB cards: train in the cloud, generate at home. A LoRA trained on a rented A100 comes back as a small safetensors file that runs fine in a quantized local pipeline — see the best Flux setup for 8GB VRAM for what inference looks like on small cards.
• Buying for this workload: a used RTX 3090 is the cheapest 24GB ticket and covers every 24GB config in the table; step up only if you also want the speed.

The cloud path is not a defeat. The config file is identical locally and rented — build the job on your laptop, debug the dataset with a cheap short run, and rent the big card only for the real training.

---

Honest Limitations {#limitations}

AI-Toolkit is the right tool for this wave, but go in with clear eyes:

• NVIDIA-first. The requirements list an NVIDIA GPU. Launcher scripts exist for macOS and Windows, but if you are not on NVIDIA silicon, treat training as experimental and expect to be your own support.
• The examples folder lags the model wave. FLUX.2 and Z-Image are supported in code but have no reference yaml yet — the web UI is the real interface for the newest models, and VRAM expectations for them are inference, not published spec. This guide will be updated when reference configs land.
• VRAM filenames are envelopes, not minimums-with-headroom. The 24GB configs achieve 24GB with quantization, batch 1, gradient checkpointing and capped buckets already applied. There is no slack hiding in them.
• No training-time promises. We deliberately quote none anywhere on this page; anyone who gives you "45 minutes on a 4090" without your dataset, resolution and step count is guessing.
• The MIT license covers the tool, not the models. Z-Image, Qwen-Image and Chroma are Apache 2.0; FLUX.2-dev is non-commercial per its model card. Check the model card before client work.
• Defaults are starting points. 2,000 steps at lr 1e-4 is where the author starts you, not a guarantee of a good LoRA — the sample images every 250 steps are the actual quality control.

---

Sources {#sources}

• ostris/ai-toolkit — README, config/examples/ configs, and extensions_built_in/diffusion_models/ source; stars/forks via the GitHub API (11,596 / 1,465, retrieved August 5, 2026)
• black-forest-labs/FLUX.2-dev — 32B parameter count and license (Hugging Face model card)
• Tongyi-MAI/Z-Image-Turbo — 6B parameter count, Apache 2.0, 16GB inference claim (Hugging Face model card)
• Qwen/Qwen-Image — 20B parameter count, Apache 2.0 (Hugging Face model card)
• black-forest-labs/FLUX.2-klein-4B — 4B parameters, Apache 2.0, ~13GB inference claim (Hugging Face model card) · lodestones/Chroma — Apache 2.0
• kohya-ss/sd-scripts — supported-architecture list (README, retrieved August 2026)

---

FAQ {#faq}