training/lora_trainer/README.md

NumZoo LoRA — training setup (reproducible)

These are the exact files used to train the NumZoo style LoRA for FLUX.2-klein-4B on Modal using ostris/ai-toolkit.

• Published LoRA: https://huggingface.co/goumsss/numzoo-flux2-klein-lora
• Base model: black-forest-labs/FLUX.2-klein-base-4B (train on base, infer on distilled)
• Trainer: ai-toolkit · Compute: Modal A100 · ~45 min · LoRA rank 32, 1500 steps

Files

File	Role
run_modal.py	Modal app (image, GPU, volume, secret) — rewritten for Modal ≥ 1.0
numzoo_klein.yaml	ai-toolkit training config (arch: flux2_klein_4b, trigger, sample prompts)

The dataset

The 54 training images + captions live one level up in ../ (image_001.jpg + image_001.txt, …). They were generated with Qwen-Image (see ../../scripts/generate_dataset.py) and captioned NUMZOO. <content> — content only, no style words, so the trigger learns the style.

How to reproduce

# 1. Clone ai-toolkit and drop these in
git clone https://github.com/ostris/ai-toolkit
cp run_modal.py        ai-toolkit/run_modal.py
cp numzoo_klein.yaml   ai-toolkit/config/numzoo_klein.yaml
cp -r <dataset>        ai-toolkit/numzoo-dataset      # the 54 image+txt pairs

# 2. Modal: install, auth, and store your HF token (needs the base-model license accepted)
pip install modal && modal token new
modal secret create huggingface HF_TOKEN=hf_xxx

# 3. Train (from inside ai-toolkit/) — checkpoints + samples land on the volume
cd ai-toolkit
modal run run_modal.py --config-file-list-str=/root/ai-toolkit/config/numzoo_klein.yaml

# 4. Download results, pick the best checkpoint (we used step 1250)
modal volume get flux-lora-models numzoo_klein_lora ./lora_output

Gotchas we hit (so you don't)

• Modal ≥ 1.0 removed modal.Mount — run_modal.py here uses Image.add_local_dir
  • add_local_file instead, and moves heavy imports inside the container function.
• pip resolution-too-deep with unpinned deps — we install from ai-toolkit's own pinned requirements.txt (which also pins the diffusers commit that supports FLUX.2), plus torch==2.7.1/torchvision/torchaudio first.
• Gated base model — accept the license for FLUX.2-klein-base-4B on HF first.
• Trigger rendered as text at inference — append no text to the prompt; the distilled 4-step klein otherwise draws the NUMZOO trigger as a literal sign.