LoliRimuru/BooruPromptGenerator

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

pip install -r requirements.txt

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

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

Generate prompts

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