README.md

---
title: NumZoo
emoji: 🦁
colorFrom: purple
colorTo: pink
sdk: gradio
sdk_version: "6.17.3"
app_file: app.py
pinned: false
license: apache-2.0
tags:
  - flux
  - text-to-image
  - kids
  - education
  - math
---

# 🦁 NumZoo

A mental math app for kids — answer questions, earn cute AI-generated animal images as rewards!

> 🏆 **Built for the [🤗 Hugging Face "Small Models, Big Adventures" hackathon](https://huggingface.co/build-small-hackathon).**
> A small (4B) image model, fine-tuned with a custom LoRA, turned into a delightful kids' reward loop.

## How it works

1. Enter your name
2. Pick your favourite animals 🐾 and places 🌍 (used to personalise reward images)
3. Answer math questions — every **3 correct answers** earns a reward image
4. Level up: **Additions → Subtractions → Multiplications → Mix**

### See it in action

| 1 · Pick animals & places | 2 · Solve math | 3 · Earn cute reward |
|:---:|:---:|:---:|
| ![picker](docs/app_picker.jpg) | ![quiz](docs/app_quiz.jpg) | ![reward](docs/app_reward.jpg) |

## Levels

| Level | Operation | Goal |
|-------|-----------|------|
| 1 | ➕ Additions | 5 correct |
| 2 | ➖ Subtractions | 5 correct |
| 3 | ✖️ Multiplications | 5 correct |
| 4 | 🎲 Mix | Endless |

## Image model

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

> Requires Python on Apple Silicon (arm64). Recommended: [Miniforge](https://github.com/conda-forge/miniforge).

```bash
# 1. Install dependencies
~/miniforge3/bin/pip install -r requirements.txt

# 2. Accept FLUX.2-klein-4B license on HuggingFace
#    → https://huggingface.co/black-forest-labs/FLUX.2-klein-4B
#    Then log in:
hf auth login

# 3. Run
~/miniforge3/bin/python3 app.py
# Opens at http://localhost:7860
# First run downloads ~23GB of model weights (cached after that)
```

> On standard Python/pip (non-Apple Silicon):
> ```bash
> pip install -r requirements.txt
> python app.py
> ```