---
license: apache-2.0
tags:
- anime
- danbooru
- prompt
---
# Booru Smart Prompt Generator

A Transformer tag language model trained on Danbooru-style metadata that generates coherent tag prompts.

## Files

- `model.safetensors` — model weights
- `config.json` — model architecture / generation defaults
- `vocab.json` — tag vocabulary
- `counts.json` — per-tag occurrence counts
- `mutex.json` — mutually exclusive tag pairs
- `model.py` — self-contained model + sampler
- `inference.py` — command-line generator

## Install

```bash
pip install -r requirements.txt
```

For the RTX 5090 / CUDA 13.x setup used during training, install the matching PyTorch wheel, e.g.:

```bash
pip install torch==2.12.1+cu130 --index-url https://download.pytorch.org/whl/cu130
```

## Generate prompts

```bash
# Empirical mode (follow Booru distribution)
python inference.py --mode empirical --count 10 --length 30

# Diverse mode (boost rare tags)
python inference.py --mode diverse --count 10 --length 30

# Anchor + blacklist + content rating
python inference.py \
  --mode empirical --count 5 --length 30 --rating s \
  --anchor "1girl,black_hair" --blacklist "1boy,smile"

# Constrained sampling (top-k / nucleus)
python inference.py \
  --mode empirical --count 5 --length 30 \
  --temperature 0.8 --top-k 50 --top-p 0.95
```

## Parameters

| Flag | Default | Description |
|------|---------|-------------|
| `--mode` | — | `empirical` (alpha=0) or `diverse` (alpha=1) |
| `--alpha` | 0.0 | Fine-grained distribution coefficient (0=empirical, 1=diverse) |
| `--rating` | `g` | Content rating token: `g` (general), `s` (sensitive), `q` (questionable), `e` (explicit) |
| `--count` | 10 | Number of prompts to generate |
| `--length` | 30 | Tags per prompt, maximum 128 |
| `--anchor` | — | Comma-separated tags that must appear in every prompt |
| `--blacklist` | — | Comma-separated tags the model must not generate |
| `--min-prob` | 0.0005 | Minimum raw model probability for a candidate tag |
| `--temperature` | 1.0 | Sampling temperature: lower = more focused, higher = more random |
| `--top-k` | 0 | Top-k sampling: keep only the k most likely tags. 0 disables it |
| `--top-p` | 1.0 | Nucleus / top-p sampling: keep the smallest set whose cumulative probability exceeds p. 1.0 disables it |
| `--distribution-weight` | 0.75 | Strength of the empirical/diverse bias applied to model scores |
| `--seed` | — | Random seed for reproducible generation |

## Notes

- Maximum prompt length is **128 tags** (including anchor), but quality is best around 30.
- Tags not in `vocab.json` are ignored for anchor/blacklist.