Upload 10 files

README.md

@@ -1,3 +1,68 @@
----
-license: apache-2.0
----
+# 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.

config.json

@@ -0,0 +1,18 @@
+{
+  "vocab_size": 25871,
+  "d_model": 256,
+  "nhead": 4,
+  "num_layers": 4,
+  "dim_feedforward": 1024,
+  "dropout": 0.1,
+  "max_len": 128,
+  "pad_id": 0,
+  "bos_id": 1,
+  "eos_id": 2,
+  "rating_g": 3,
+  "rating_s": 4,
+  "rating_q": 5,
+  "rating_e": 6,
+  "offset": 7,
+  "distribution_weight": 0.5
+}

inference.py

@@ -0,0 +1,100 @@
+"""Standalone inference script for the Booru prompt generator release."""
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_model
+
+from model import Vocab, SimpleGraph, TagTransformer, NeuralPromptGenerator
+
+
+def load_generator(release_dir: str, seed: Optional[int] = None):
+    release_dir = Path(release_dir)
+    with open(release_dir / "config.json", "r", encoding="utf-8") as f:
+        config = json.load(f)
+    with open(release_dir / "vocab.json", "r", encoding="utf-8") as f:
+        tags = json.load(f)
+    with open(release_dir / "counts.json", "r", encoding="utf-8") as f:
+        counts = json.load(f)
+    with open(release_dir / "mutex.json", "r", encoding="utf-8") as f:
+        mutex = json.load(f)
+
+    vocab = Vocab(tags, counts)
+    graph = SimpleGraph(vocab, mutex)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = TagTransformer(
+        vocab_size=len(vocab),
+        d_model=config["d_model"],
+        nhead=config["nhead"],
+        num_layers=config["num_layers"],
+        dim_feedforward=config["dim_feedforward"],
+        dropout=config["dropout"],
+        max_len=config["max_len"] + 2,
+    ).to(device)
+    load_model(model, str(release_dir / "model.safetensors"))
+
+    return NeuralPromptGenerator(
+        model,
+        graph,
+        device=device,
+        seed=seed,
+        distribution_weight=config.get("distribution_weight", 0.75),
+    )
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate Booru tag prompts.")
+    parser.add_argument("--release-dir", default=".", help="Path to the release folder")
+    parser.add_argument("--mode", choices=["empirical", "diverse"], default=None)
+    parser.add_argument("--alpha", type=float, default=None)
+    parser.add_argument("--count", type=int, default=10)
+    parser.add_argument("--length", type=int, default=30)
+    parser.add_argument("--anchor", default="", help="Comma-separated anchor tags")
+    parser.add_argument("--blacklist", default="", help="Comma-separated tags to forbid")
+    parser.add_argument("--rating", default="g", choices=["g", "s", "q", "e"],
+                        help="Content rating token to condition on")
+    parser.add_argument("--min-prob", type=float, default=0.0005)
+    parser.add_argument("--temperature", type=float, default=1.0)
+    parser.add_argument("--top-k", type=int, default=0, help="Top-k sampling (0 = disabled)")
+    parser.add_argument("--top-p", type=float, default=1.0, help="Nucleus/top-p sampling (1.0 = disabled)")
+    parser.add_argument("--distribution-weight", type=float, default=None)
+    parser.add_argument("--seed", type=int, default=None)
+    args = parser.parse_args()
+
+    alpha = args.alpha
+    if args.mode == "empirical":
+        alpha = 0.0
+    elif args.mode == "diverse":
+        alpha = 1.0
+    if alpha is None:
+        alpha = 0.0
+
+    gen = load_generator(args.release_dir, seed=args.seed)
+    if args.distribution_weight is not None:
+        gen.distribution_weight = args.distribution_weight
+
+    anchor_tags = [t.strip() for t in args.anchor.split(",") if t.strip()]
+    blacklist_tags = [t.strip() for t in args.blacklist.split(",") if t.strip()]
+
+    prompts = gen.generate(
+        alpha=alpha,
+        count=args.count,
+        length=args.length,
+        anchor=anchor_tags or None,
+        blacklist=blacklist_tags or None,
+        min_prob=args.min_prob,
+        temperature=args.temperature,
+        top_k=getattr(args, "top_k", 0),
+        top_p=getattr(args, "top_p", 1.0),
+        rating=args.rating,
+    )
+
+    for tags in prompts:
+        print(", ".join(tags))
+
+
+if __name__ == "__main__":
+    main()

model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:df1d367f5d3a6fa475f05ac24c0716fb256c987e1ab003350ecb3c52331cbf06
+size 65754008

model.py

@@ -0,0 +1,277 @@
+"""Minimal self-contained inference model for the Booru prompt generator release."""
+import math
+import re
+from typing import Dict, List, Optional, Set
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+PAD_ID = 0
+BOS_ID = 1
+EOS_ID = 2
+RATING_G = 3
+RATING_S = 4
+RATING_Q = 5
+RATING_E = 6
+OFFSET = 7
+
+_RATING_TOKENS = {"g": RATING_G, "s": RATING_S, "q": RATING_Q, "e": RATING_E}
+
+
+def _apply_top_k_top_p(probs: torch.Tensor, top_k: int, top_p: float) -> torch.Tensor:
+    """Filter a probability distribution with top-k and/or nucleus (top-p) sampling."""
+    if top_k > 0:
+        k = min(top_k, probs.size(0))
+        threshold = torch.topk(probs, k).values[-1]
+        probs = probs.where(probs >= threshold, torch.zeros_like(probs))
+    if top_p < 1.0:
+        sorted_probs, sorted_idx = torch.sort(probs, descending=True)
+        cumsum = torch.cumsum(sorted_probs, dim=0)
+        nucleus_mask = cumsum <= top_p
+        if nucleus_mask.any():
+            nucleus_mask[0] = True
+        kept = torch.zeros_like(probs, dtype=torch.bool)
+        kept.scatter_(0, sorted_idx, nucleus_mask)
+        probs = probs.where(kept, torch.zeros_like(probs))
+    return probs
+
+
+class Vocab:
+    def __init__(self, tags: List[str], counts: List[int]):
+        self.tags = tags
+        self.tag_to_idx = {tag: i for i, tag in enumerate(tags)}
+        self.counts = np.array(counts, dtype=np.int64)
+        self.total = int(self.counts.sum())
+        self.freqs = self.counts.astype(np.float64) / max(self.total, 1)
+
+    def __len__(self) -> int:
+        return len(self.tags)
+
+
+class SimpleGraph:
+    def __init__(self, vocab: Vocab, mutex: List[List[int]]):
+        self.vocab = vocab
+        self.mutex = mutex
+
+
+class TagTransformer(nn.Module):
+    """Permutation-equivariant Transformer for tag-set generation."""
+
+    def __init__(
+        self,
+        vocab_size: int,
+        d_model: int = 256,
+        nhead: int = 4,
+        num_layers: int = 4,
+        dim_feedforward: int = 1024,
+        dropout: float = 0.1,
+        max_len: int = 256,
+    ):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.max_len = max_len
+        self.embedding = nn.Embedding(vocab_size + OFFSET, d_model, padding_idx=PAD_ID)
+        layer = nn.TransformerEncoderLayer(
+            d_model=d_model,
+            nhead=nhead,
+            dim_feedforward=dim_feedforward,
+            dropout=dropout,
+            batch_first=True,
+            norm_first=True,
+        )
+        self.transformer = nn.TransformerEncoder(layer, num_layers=num_layers, enable_nested_tensor=False)
+        self.output = nn.Linear(d_model, vocab_size + OFFSET)
+        self._init_weights()
+
+    def _init_weights(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        seq_len = input_ids.size(1)
+        mask = nn.Transformer.generate_square_subsequent_mask(seq_len, device=input_ids.device).bool()
+        x = self.embedding(input_ids)
+        padding_mask = input_ids == PAD_ID
+        x = self.transformer(
+            x,
+            mask=mask,
+            src_key_padding_mask=padding_mask,
+            is_causal=True,
+        )
+        return self.output(x)
+
+
+class NeuralPromptGenerator:
+    GROUP_SUFFIXES = ["_hair", "_eyes"]
+
+    def __init__(
+        self,
+        model: TagTransformer,
+        graph: SimpleGraph,
+        device: Optional[torch.device] = None,
+        seed: Optional[int] = None,
+        distribution_weight: float = 0.75,
+        fallback_top_k: int = 20,
+    ):
+        self.model = model
+        self.graph = graph
+        self.vocab = graph.vocab
+        self.device = device or torch.device("cpu")
+        self.model.to(self.device)
+        self.model.eval()
+        self.rng = np.random.default_rng(seed)
+        self.distribution_weight = distribution_weight
+        self.fallback_top_k = fallback_top_k
+
+        self._n_real = len(self.vocab)
+        self._real_token_ids = torch.arange(self._n_real, device=self.device) + OFFSET
+        self._log_uniform = -np.log(max(self._n_real, 1))
+        self._special_ids = {PAD_ID, BOS_ID, EOS_ID}
+
+        self._group_ids = np.full(self._n_real, -1, dtype=np.int32)
+        for i, tag in enumerate(self.vocab.tags):
+            for gid, suffix in enumerate(self.GROUP_SUFFIXES):
+                if tag.endswith(suffix):
+                    self._group_ids[i] = gid
+                    break
+        self._group_ids_tensor = torch.from_numpy(self._group_ids).to(self.device)
+
+    def _target_log_prob(self, idx: int, alpha: float) -> float:
+        log_emp = np.log(max(self.vocab.freqs[idx], 1e-12))
+        return (1.0 - 2.0 * alpha) * log_emp
+
+    def _num_people(self, tag_idxs: Set[int]) -> int:
+        n = 0
+        for idx in tag_idxs:
+            tag = self.vocab.tags[idx]
+            if tag == "solo":
+                n = max(n, 1)
+            elif tag in ("1girl", "1boy", "1other"):
+                n = max(n, 1)
+            elif tag in ("2girls", "2boys", "2others"):
+                n = max(n, 2)
+            elif tag in ("3girls", "3boys", "3others"):
+                n = max(n, 3)
+            elif tag in ("4girls", "4boys", "4others"):
+                n = max(n, 4)
+            elif tag in ("5girls", "5boys", "5others"):
+                n = max(n, 5)
+            elif tag in (
+                "6+girls", "6+boys", "6+others",
+                "multiple_girls", "multiple_boys", "multiple_others",
+            ):
+                n = max(n, 2)
+        return max(n, 1)
+
+    def generate(
+        self,
+        alpha: float,
+        count: int,
+        length: int = 30,
+        anchor: Optional[List[str]] = None,
+        blacklist: Optional[List[str]] = None,
+        min_prob: float = 0.0,
+        temperature: float = 1.0,
+        top_k: int = 0,
+        top_p: float = 1.0,
+        rating: Optional[str] = None,
+    ) -> List[List[str]]:
+        anchor = anchor or []
+        anchor_indices = [self.vocab.tag_to_idx[t] for t in anchor if t in self.vocab.tag_to_idx]
+        blacklist = blacklist or []
+        blacklist_indices = {self.vocab.tag_to_idx[t] for t in blacklist if t in self.vocab.tag_to_idx}
+        rating_token = _RATING_TOKENS.get((rating or "g").lower(), RATING_G)
+
+        target_bias = torch.tensor(
+            [self._target_log_prob(i, alpha) for i in range(self._n_real)],
+            dtype=torch.float32,
+            device=self.device,
+        )
+
+        results: List[List[str]] = []
+        with torch.no_grad():
+            for _ in range(count):
+                prompt_tokens: List[int] = [BOS_ID, rating_token] + [idx + OFFSET for idx in anchor_indices]
+                present_tag_idxs: Set[int] = set(anchor_indices)
+                excluded_tag_idxs: Set[int] = set(blacklist_indices)
+                group_counts: Dict[int, int] = {}
+                for idx in anchor_indices:
+                    excluded_tag_idxs.update(self.graph.mutex[idx])
+                    gid = self._group_ids[idx]
+                    if gid >= 0:
+                        group_counts[gid] = group_counts.get(gid, 0) + 1
+
+                target_len = max(length, len(anchor_indices))
+                max_len = getattr(self.model, "max_len", 256)
+                if target_len > max_len - 2:
+                    target_len = max_len - 2
+
+                while len(prompt_tokens) - 1 < target_len:
+                    input_ids = torch.tensor([prompt_tokens], dtype=torch.long, device=self.device)
+                    logits = self.model(input_ids)[:, -1, :]
+
+                    model_probs = torch.softmax(logits / max(temperature, 1e-6), dim=-1).squeeze(0)
+                    allowed = model_probs >= min_prob
+
+                    real_allowed = allowed[self._real_token_ids]
+                    if not real_allowed.any():
+                        k = min(self.fallback_top_k, self._n_real)
+                        topk = torch.topk(model_probs[self._real_token_ids], k=k).indices
+                        real_allowed = torch.zeros(self._n_real, dtype=torch.bool, device=self.device)
+                        real_allowed[topk] = True
+                        allowed = allowed.clone()
+                        allowed[self._real_token_ids] = real_allowed
+
+                    max_people = self._num_people(present_tag_idxs)
+                    full_group_ids = [gid for gid, c in group_counts.items() if c >= max_people]
+
+                    biased_logits = logits.squeeze(0).clone()
+                    biased_logits[self._real_token_ids] += self.distribution_weight * target_bias
+
+                    biased_logits[~allowed] = -float("inf")
+                    for tid in self._special_ids:
+                        biased_logits[tid] = -float("inf")
+                    for idx in present_tag_idxs:
+                        biased_logits[idx + OFFSET] = -float("inf")
+                    for idx in excluded_tag_idxs:
+                        biased_logits[idx + OFFSET] = -float("inf")
+                    if full_group_ids:
+                        full_groups_tensor = torch.tensor(full_group_ids, dtype=torch.int32, device=self.device)
+                        group_full_mask = torch.isin(self._group_ids_tensor, full_groups_tensor)
+                        biased_logits[self._real_token_ids[group_full_mask]] = -float("inf")
+
+                    probs = torch.softmax(biased_logits / max(temperature, 1e-6), dim=0)
+                    probs = _apply_top_k_top_p(probs, top_k, top_p)
+                    if not torch.isfinite(probs).all() or probs.sum() <= 0:
+                        break
+                    probs = probs / probs.sum()
+                    probs = probs.cpu().numpy()
+                    probs = probs / probs.sum()
+
+                    token_id = int(self.rng.choice(len(probs), p=probs))
+                    if token_id in self._special_ids:
+                        break
+                    tag_idx = token_id - OFFSET
+                    if tag_idx < 0 or tag_idx >= self._n_real:
+                        break
+                    if tag_idx in present_tag_idxs or tag_idx in excluded_tag_idxs:
+                        break
+
+                    prompt_tokens.append(token_id)
+                    present_tag_idxs.add(tag_idx)
+                    excluded_tag_idxs.update(self.graph.mutex[tag_idx])
+                    gid = self._group_ids[tag_idx]
+                    if gid >= 0:
+                        group_counts[gid] = group_counts.get(gid, 0) + 1
+
+                results.append([
+                    self.vocab.tags[i - OFFSET]
+                    for i in prompt_tokens
+                    if i >= OFFSET
+                ])
+
+        return results

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:79fa338fbdd18c502f462b43117109b279775a17bf25b3764fdbb56a094f62a1
+size 65743184

requirements.txt

@@ -0,0 +1,3 @@
+torch>=2.12
+numpy>=1.26
+safetensors>=0.4