# NumZoo LoRA — training setup (reproducible)

These are the exact files used to train the **NumZoo style LoRA** for FLUX.2-klein-4B
on [Modal](https://modal.com) using [ostris/ai-toolkit](https://github.com/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`](../../scripts/generate_dataset.py)) and captioned
`NUMZOO. <content>` — content only, no style words, so the trigger learns the style.

## How to reproduce
```bash
# 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.