Docs: publish LoRA trainer + add README fine-tuning story

Task #3 — training/lora_trainer/: the reproducible Modal + ai-toolkit setup
(run_modal.py, numzoo_klein.yaml, README with repro steps + the gotchas we hit).

Task #4 — README: new "Custom AI art: the NumZoo LoRA" section with before/after
panda (photoreal → cozy), a LoRA gallery, the Qwen-Image dataset story, training
summary, and a link to the published HF LoRA (goumsss/numzoo-flux2-klein-lora).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

README.md

@@ -38,8 +38,42 @@ A mental math app for kids — answer questions, earn cute AI-generated animal i
 
 ## Image model
 
-Rewards generated with **FLUX.2-klein-4B** (4B params, Apache 2.0) via 🤗 Diffusers.  
-Images start generating in the background as soon as the quiz begins.
+Rewards are generated with **FLUX.2-klein-4B** (4B params, Apache 2.0) via 🤗 Diffusers,
+plus a **custom NumZoo style LoRA** (below). Images start generating in the background as
+soon as the quiz begins, so a reward is usually ready the moment it's earned.
+
+## ✨ Custom AI art: the NumZoo LoRA
+
+Out of the box, FLUX.2-klein renders the same prompt in wildly different styles — often
+photorealistic — which doesn't fit a soft, cozy kids' app. So we fine-tuned a **style
+LoRA** that pins every reward to the same kawaii children's-book look.
+
+**Before → after** (same prompt: *"a cute baby panda on a snowy mountain top"*):
+
+| Base FLUX.2-klein-4B | + NumZoo LoRA |
+|:---:|:---:|
+| ![before](docs/panda_before.jpg) | ![after](docs/panda_after.jpg) |
+| photorealistic, inconsistent | cozy, on-brand, every time |
+
+More rewards from the LoRA:
+
+| | | |
+|:---:|:---:|:---:|
+| ![bunny](docs/lora_bunny.jpg) | ![unicorn](docs/lora_unicorn.jpg) | ![fox](docs/lora_fox.jpg) |
+
+### How we made it
+1. **Dataset** — 54 cozy scenes generated with **Qwen-Image** (12 animals × 10 places,
+   incl. multi-animal/multi-place combos), captioned `NUMZOO. <content>` with the style
+   left *undescribed* so the trigger word carries it. See
+   [`scripts/generate_dataset.py`](scripts/generate_dataset.py) and
+   [`training/`](training/) for the images + captions.
+2. **Training** — LoRA (rank 32, 1500 steps) on `FLUX.2-klein-base-4B` via
+   [ostris/ai-toolkit](https://github.com/ostris/ai-toolkit) on a **Modal** A100 (~45 min).
+   Reproducible setup in [`training/lora_trainer/`](training/lora_trainer/).
+3. **Inference** — the LoRA loads on the distilled 4-step klein in `image_generator.py`;
+   the app simply prepends the `NUMZOO` trigger to every prompt.
+
+**Published LoRA:** 🤗 [goumsss/numzoo-flux2-klein-lora](https://huggingface.co/goumsss/numzoo-flux2-klein-lora)
 
 ## Run locally
 

training/lora_trainer/README.md

@@ -0,0 +1,50 @@
+# 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.

training/lora_trainer/numzoo_klein.yaml

@@ -0,0 +1,69 @@
+---
+# NumZoo style LoRA — FLUX.2-klein-4B, trained on Modal via ai-toolkit.
+# Trained on the BASE model (FLUX.2-klein-base-4B); the resulting LoRA is loaded
+# on the DISTILLED FLUX.2-klein-4B at inference in the app (BFL-recommended).
+job: extension
+config:
+  name: "numzoo_klein_lora"          # output folder/filename name
+  process:
+    - type: 'sd_trainer'
+      training_folder: "/root/ai-toolkit/modal_output"   # must match MOUNT_DIR in run_modal.py
+      device: cuda:0
+      # Captions already start with "NUMZOO." so we do NOT set trigger_word here
+      # (avoids double-injecting it). Sample prompts below include NUMZOO explicitly.
+      network:
+        type: "lora"
+        linear: 32
+        linear_alpha: 32
+      save:
+        dtype: float16
+        save_every: 250            # checkpoints at 250,500,...,1500 — pick the best
+        max_step_saves_to_keep: 6
+      datasets:
+        - folder_path: "/root/ai-toolkit/numzoo-dataset"
+          caption_ext: "txt"
+          caption_dropout_rate: 0.05
+          shuffle_tokens: false
+          cache_latents_to_disk: true
+          resolution: [ 512, 768, 1024 ]
+      train:
+        batch_size: 1
+        steps: 1500                 # 54 imgs; style LoRA peaks ~750–1500
+        gradient_accumulation_steps: 1
+        train_unet: true
+        train_text_encoder: false
+        gradient_checkpointing: true
+        noise_scheduler: "flowmatch"
+        timestep_type: "weighted"   # klein-recommended (UI default)
+        optimizer: "adamw8bit"
+        lr: 1e-4
+        ema_config:
+          use_ema: true
+          ema_decay: 0.99
+        dtype: bf16
+      model:
+        name_or_path: "black-forest-labs/FLUX.2-klein-base-4B"
+        arch: "flux2_klein_4b"      # resolves to Flux2Klein4BModel (TE=Qwen3-4B, VAE auto)
+        quantize: true              # 8-bit to fit comfortably on the GPU
+        quantize_te: true
+        qtype: "qfloat8"
+        low_vram: true
+      sample:
+        sampler: "flowmatch"
+        sample_every: 250
+        width: 512                  # match the app's reward image size
+        height: 512
+        guidance_scale: 4           # base (non-distilled) model uses real CFG
+        sample_steps: 25            # base model needs more steps than distilled
+        seed: 42
+        walk_seed: true
+        prompts:
+          - "NUMZOO. A cute bunny in an enchanted mushroom forest"
+          - "NUMZOO. A cute baby panda on a snowy mountain top"
+          - "NUMZOO. A cute puppy and kitten in a cosy cottage garden"
+          - "NUMZOO. A cute baby lion, baby tiger and bunny surrounded by sparkling stars"
+          - "NUMZOO. A cute unicorn under a rainbow and on a glowing crescent moon"
+          - "NUMZOO. A cute baby fox on a sunny beach with ocean waves"
+meta:
+  name: "[name]"
+  version: '1.0'

training/lora_trainer/run_modal.py

@@ -0,0 +1,90 @@
+'''
+ostris/ai-toolkit on https://modal.com  —  rewritten for Modal >= 1.0 API.
+
+Run training (from inside the ai-toolkit directory) with:
+    modal run run_modal.py --config-file-list-str=/root/ai-toolkit/config/numzoo_klein.yaml
+'''
+
+import os
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+os.environ["DISABLE_TELEMETRY"] = "YES"
+import modal
+
+# --- paths -----------------------------------------------------------------
+LOCAL_AI_TOOLKIT = os.path.dirname(os.path.abspath(__file__))  # this repo, shipped to the container
+REMOTE_AI_TOOLKIT = "/root/ai-toolkit"
+MOUNT_DIR = "/root/ai-toolkit/modal_output"  # volume mount (must be an empty path)
+
+# --- persistent volume for outputs (LoRA weights, samples) -----------------
+model_volume = modal.Volume.from_name("flux-lora-models", create_if_missing=True)
+
+# --- container image: ai-toolkit deps + the repo + dataset shipped in ------
+# Use ai-toolkit's OWN pinned requirements (requirements.txt -> requirements_base.txt).
+# They pin a specific diffusers commit that supports FLUX.2 klein and avoid the pip
+# "resolution-too-deep" blowup that unpinned deps (esp. old controlnet_aux) caused.
+image = (
+    modal.Image.debian_slim(python_version="3.11")
+    .apt_install("libgl1", "libglib2.0-0", "git")
+    # CUDA-enabled torch first, pinned to what torchao/torchcodec expect.
+    # torchaudio is imported by ai-toolkit's config_modules — must match torch ver.
+    .pip_install("torch==2.7.1", "torchvision==0.22.1", "torchaudio==2.7.1")
+    # copy both requirements files into the build context, then install (pinned)
+    .add_local_file(os.path.join(LOCAL_AI_TOOLKIT, "requirements.txt"),
+                    "/tmp/reqs/requirements.txt", copy=True)
+    .add_local_file(os.path.join(LOCAL_AI_TOOLKIT, "requirements_base.txt"),
+                    "/tmp/reqs/requirements_base.txt", copy=True)
+    .run_commands("cd /tmp/reqs && pip install -r requirements.txt")
+    # ship the ai-toolkit repo (code + numzoo-dataset/) into the container LAST
+    # (copy=False runtime mount must be the final image layer).
+    .add_local_dir(
+        LOCAL_AI_TOOLKIT,
+        REMOTE_AI_TOOLKIT,
+        ignore=[
+            ".git", "**/.git", "**/node_modules", "**/__pycache__", "**/*.pyc",
+            "modal_output", "**/.venv", "venv", "ui/node_modules", "output",
+        ],
+    )
+)
+
+app = modal.App(name="numzoo-lora-training", image=image, volumes={MOUNT_DIR: model_volume})
+
+
+@app.function(
+    gpu="A100",                 # 40GB — plenty for a quantized 4B LoRA (<24GB)
+    timeout=7200,               # 2h ceiling; a 1500-step run is ~45 min
+    secrets=[modal.Secret.from_name("huggingface")],  # provides HF_TOKEN for gated model
+)
+def main(config_file_list_str: str, recover: bool = False, name: str = None):
+    import sys
+    sys.path.insert(0, REMOTE_AI_TOOLKIT)
+    from toolkit.job import get_job  # imported in-container (pulls torch etc.)
+
+    config_file_list = config_file_list_str.split(",")
+    jobs_completed = 0
+    jobs_failed = 0
+    print(f"Running {len(config_file_list)} job{'' if len(config_file_list) == 1 else 's'}")
+
+    for config_file in config_file_list:
+        try:
+            job = get_job(config_file, name)
+            job.config['process'][0]['training_folder'] = MOUNT_DIR
+            os.makedirs(MOUNT_DIR, exist_ok=True)
+            print(f"Training outputs will be saved to: {MOUNT_DIR}")
+            job.run()
+            model_volume.commit()   # persist weights/samples after each job
+            job.cleanup()
+            jobs_completed += 1
+        except Exception as e:
+            print(f"Error running job: {e}")
+            jobs_failed += 1
+            if not recover:
+                raise
+
+    print("========================================")
+    print(f"Result: {jobs_completed} completed, {jobs_failed} failed")
+    print("========================================")
+
+
+@app.local_entrypoint()
+def run(config_file_list_str: str, recover: bool = False, name: str = None):
+    main.remote(config_file_list_str=config_file_list_str, recover=recover, name=name)