File size: 2,639 Bytes
fb57054 8b27df3 fb57054 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | ---
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. |