LoliRimuru/BooruPromptGenerator

🤗 Hugging Face 来源apache-2.016M 参数66 MBsafetensors✓ 2 个校验和今天更新
需要做种者 →

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.