Training

Training

brandwright writes plain ostris ai-toolkit JSON configs and can start them. The same file runs by hand with python run.py config.json from the ai-toolkit folder.

Base models (checked 2026-10-15)

--model	Licence	Min VRAM	Notes
zimage (default)	Apache-2.0	16 GB	strong text, any aspect 512-2048, residual site LoRAs
qwen_image (2512)	Apache-2.0	24 GB	best open text and layout; 3-bit with an accuracy recovery adapter
klein4b (FLUX.2 klein base 4B)	Apache-2.0	12 GB	fastest, weaker text
zimage_turbo	Apache-2.0	16 GB	trains with ostris's turbo adapter, generates in 9 steps; no residual mode

Not for company work:

• Qwen-Image-2.1: research licence.
• FLUX.2 dev, klein 9B and FLUX.1 dev: non-commercial.
• Ideogram 4: non-commercial.

Two have revenue limits: FLUX 3 klein (under US$5M) and Krea 2 (under US$1M).

Brand first, then each site

BW train KIT --lora brand --run
BW status KIT
BW pick KIT --lora brand --step 300 --to <ComfyUI>\models\loras
BW train KIT --lora portal --run           # residual on Z-Image, continue elsewhere; --mode to choose

Defaults:

• rank 32 and lr 1e-4;
• buckets 768 and 1024;
• about 5 passes over the pictures (5 steps a picture, between 400 and 1,500);
• a checkpoint every 100 steps;
• samples: 4 prompts at every checkpoint for the brand LoRA; 2 prompts every 200 steps, with no step-0 sample, for a residual site LoRA (--samples, --sample-every; --samples 0 turns them off). Before 0.1.2, every LoRA was sampled like the brand LoRA.

The samples use held-out prompts (pricing table, settings page, modal, toast). The capture rarely holds those kinds, so the samples show what the LoRA generalises. The first run downloads the base model; Z-Image is about 20 GB.

Picking the checkpoint

Open train/out/<name>/samples/ at each step:

• Too weak: generic UI and wrong colours. Keep training.
• Right: the brand's colours, type and spacing on components it never saw, with readable text.
• Too far: it copies real pages, draws the same layout whatever the prompt, or the text goes soft and pages turn sparse. Use an earlier checkpoint. In the first real run (90 pictures on Z-Image), step 250 was best, and from step 500 the samples had collapsed into blurred, near-empty layouts.

When ComfyUI is running on Windows and has the old LoRA loaded, it holds the file open, so pick --to cannot replace it. It writes a copy named after the step instead (acme-brand-s250.safetensors) and points brand.json at it, so gen --lora brand uses the new one.

For a site LoRA in residual mode, the trainer draws its samples without the brand LoRA, so they look weaker than the real result. Copy the site checkpoints into ComfyUI under their own names and compare them with gen or sweep on brand + site instead. A sample takes about 2 minutes on a 16 GB card. In the second test run, 5 sets of 4 took 40 of the site LoRA's 121 minutes, which is why residual runs now sample less.

Training again

ai-toolkit continues from the newest checkpoint in train/out/<name>/. So when a LoRA is trained again with --run, bw.py first moves that folder to train/out/<name>-run-<date>, and the new run starts at step 0. Add --resume to carry on the old run instead, for example after a crash.

Generating with the LoRA

ComfyUI needs the files for the graph:

Graph	Diffusion model	Text encoder (CLIPLoader type)	VAE
zimage	z_image_bf16.safetensors	qwen_3_4b.safetensors (lumina2)	ae.safetensors
zimage_turbo	z_image_turbo_bf16.safetensors	qwen_3_4b.safetensors (lumina2)	ae.safetensors
klein4b, klein4b_base	flux-2-klein-4b.safetensors, flux-2-klein-base-4b.safetensors	qwen_3_4b.safetensors (flux2)	flux2-vae.safetensors
qwen_image	qwen_image_2512_fp8_e4m3fn.safetensors	qwen_2.5_vl_7b_fp8_scaled.safetensors (qwen_image)	qwen_image_vae.safetensors

Different file names: set them in brand.json under comfy_models (Kit layout).

Troubleshooting

• Blackwell (RTX 50xx) needs PyTorch 2.7 or later built for CUDA 12.8 or later. ai-toolkit's cu130 wheels need driver 580 or later. "no kernel image is available" means an older wheel.
• Out of memory: use --res 768 or --res 512 768, keep the rank at 32 or lower, and make sure ComfyUI or another job is not holding the GPU.
• The log is not moving: the first run downloads the model and caches latents and text embeddings. BW status prints the last log line.
• No GPU here: write the config without --run and run it on a training machine, or on a private GPU in your company's cloud.

Other trainers

musubi-tuner, OneTrainer, SimpleTuner and diffusion-pipe all read the same dataset folders (a PNG plus a same-name .txt). musubi's --base_weights gives the same differential training as residual mode.

Since 0.1.3, bw.py train KIT --lora LORA --for musubi writes the same job for musubi-tuner: train/<name>-musubi/dataset.toml and a train.sh with the cache and training commands (set the DiT, VAE and text-encoder paths at its top). Residual site LoRAs use --base_weights. --for onetrainer writes config.json and concepts.json for python scripts/train.py --config-path. OneTrainer trains by epochs and cannot merge a LoRA during training, so it refuses residual mode; use --mode continue. The formats come from musubi-tuner v0.3.6 and OneTrainer's source of 2026-10-11. Since 0.1.5 the musubi job is proven: the Z-Image brand LoRA job ran on musubi-tuner v0.3.6 with only the three model paths set, and the LoRA, converted with musubi's convert_lora.py --target other, drew on brand in ComfyUI as the ai-toolkit LoRA does. Run train.sh with musubi's venv active, since it calls a bare python. The OneTrainer job, and musubi's residual mode, have not run yet.