| --- |
| 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. |
