Upload 10 files

README.md

@@ -1,14 +1,39 @@
----
-title: BooruPromptRNG
-emoji: 📚
-colorFrom: blue
-colorTo: pink
-sdk: gradio
-sdk_version: 6.19.0
-python_version: '3.13'
-app_file: app.py
-pinned: false
-license: apache-2.0
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Booru Smart Prompt Generator — Gradio UI
+
+A web UI for the Booru tag prompt generator.
+
+## Files
+
+- `app.py` — Gradio application
+- `model.py`, `model.safetensors`, `config.json`, `vocab.json`, `counts.json`, `mutex.json` — model assets
+- `requirements.txt` — Python dependencies
+
+## Usage on Hugging Face Spaces
+
+1. Create a new **Gradio** Space.
+2. Upload the contents of this folder (including `model.safetensors`).
+3. Wait for the container to build and start.
+4. Open the app and generate prompts.
+
+## Parameters
+
+- **Mode**: `Empirical` (alpha=0), `Diverse` (alpha=1), or `Custom alpha`.
+- **Alpha**: distribution coefficient (active only in Custom mode).
+- **Tags per prompt**: 1–128; 30 is recommended.
+- **Anchor tags**: fixed starting tags the generator must include.
+- **Blacklist tags**: tags the generator must never emit.
+- **Min model probability**: rejects low-probability / conflicting candidates.
+- **Temperature**: higher = more random.
+- **Distribution bias weight**: strength of the empirical ↔ diverse reweighting.
+
+## Local run
+
+```bash
+pip install -r requirements.txt
+python app.py
+```
+
+> For an RTX 5090 / CUDA 13.x setup, install the matching PyTorch wheel before running, e.g.:
+> ```bash
+> pip install torch==2.12.1+cu130 --index-url https://download.pytorch.org/whl/cu130
+> ```

app.py

@@ -0,0 +1,124 @@
+"""Gradio app for the Booru prompt generator."""
+import json
+from pathlib import Path
+
+import gradio as gr
+import torch
+from safetensors.torch import load_model
+
+from model import Vocab, SimpleGraph, TagTransformer, NeuralPromptGenerator
+
+
+_RELEASE_DIR = Path(__file__).parent
+
+
+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"] + 1,
+).to(_DEVICE)
+load_model(_MODEL, str(_RELEASE_DIR / "model.safetensors"))
+
+
+def _parse_comma(text: str) -> list:
+    return [t.strip() for t in text.split(",") if t.strip()]
+
+
+def generate(
+    mode: str,
+    alpha: float,
+    count: int,
+    length: int,
+    anchor: str,
+    blacklist: str,
+    min_prob: float,
+    temperature: float,
+    distribution_weight: float,
+    seed: int,
+):
+    if mode == "Empirical":
+        alpha = 0.0
+    elif mode == "Diverse":
+        alpha = 1.0
+
+    gen = NeuralPromptGenerator(
+        _MODEL,
+        _GRAPH,
+        device=_DEVICE,
+        seed=None if seed < 0 else seed,
+        distribution_weight=distribution_weight,
+    )
+
+    prompts = gen.generate(
+        alpha=alpha,
+        count=count,
+        length=length,
+        anchor=_parse_comma(anchor) or None,
+        blacklist=_parse_comma(blacklist) or None,
+        min_prob=min_prob,
+        temperature=temperature,
+    )
+    return "\n".join(", ".join(p) for p in prompts)
+
+
+with gr.Blocks(title="Booru Smart Prompt Generator") as demo:
+    gr.Markdown("# Booru Smart Prompt Generator")
+    gr.Markdown(
+        "Generate semantically coherent Danbooru-style tag prompts. "
+        "Empirical mode follows the dataset distribution; Diverse mode boosts rare tags."
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            mode = gr.Dropdown(
+                choices=["Empirical", "Diverse", "Custom alpha"],
+                value="Empirical",
+                label="Mode",
+            )
+            alpha = gr.Slider(0.0, 1.0, value=0.0, step=0.05, label="Alpha (custom)")
+            count = gr.Number(value=5, precision=0, label="Number of prompts", minimum=1, maximum=100)
+            length = gr.Slider(1, 128, value=30, step=1, label="Tags per prompt")
+            anchor = gr.Textbox(
+                label="Anchor tags (comma-separated)",
+                placeholder="e.g. 1girl, black_hair",
+            )
+            blacklist = gr.Textbox(
+                label="Blacklist tags (comma-separated)",
+                placeholder="e.g. 1boy, smile",
+            )
+            min_prob = gr.Slider(0.0, 0.1, value=0.0, step=0.001, label="Min model probability threshold")
+            temperature = gr.Slider(0.1, 2.0, value=1.0, step=0.05, label="Temperature")
+            distribution_weight = gr.Slider(0.0, 2.0, value=0.75, step=0.05, label="Distribution bias weight")
+            seed = gr.Number(value=-1, precision=0, label="Seed (-1 for random)")
+            generate_btn = gr.Button("Generate", variant="primary")
+
+        with gr.Column(scale=2):
+            output = gr.Textbox(label="Generated prompts", lines=20)
+
+    generate_btn.click(
+        fn=generate,
+        inputs=[mode, alpha, count, length, anchor, blacklist, min_prob, temperature, distribution_weight, seed],
+        outputs=output,
+    )
+
+
+if __name__ == "__main__":
+    demo.launch()

config.json

@@ -0,0 +1,10 @@
+{
+  "vocab_size": 25871,
+  "d_model": 256,
+  "nhead": 4,
+  "num_layers": 4,
+  "dim_feedforward": 1024,
+  "dropout": 0.1,
+  "max_len": 128,
+  "distribution_weight": 0.75
+}

inference.py

@@ -0,0 +1,93 @@
+"""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"] + 1,
+    ).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("--min-prob", type=float, default=0.0)
+    parser.add_argument("--temperature", type=float, default=1.0)
+    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,
+    )
+
+    for tags in prompts:
+        print(", ".join(tags))
+
+
+if __name__ == "__main__":
+    main()

model.py

@@ -0,0 +1,251 @@
+"""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
+OFFSET = 3
+
+
+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 SinusoidalPositionalEncoding(nn.Module):
+    def __init__(self, d_model: int, max_len: int = 512):
+        super().__init__()
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, d_model, 2, dtype=torch.float) * (-math.log(10000.0) / d_model)
+        )
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        self.register_buffer("pe", pe.unsqueeze(0))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x + self.pe[:, : x.size(1)]
+
+
+class TagTransformer(nn.Module):
+    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.embedding = nn.Embedding(vocab_size + OFFSET, d_model, padding_idx=PAD_ID)
+        self.pos_encoding = SinusoidalPositionalEncoding(d_model, max_len=max_len)
+        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)
+        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)
+        x = self.embedding(input_ids) * math.sqrt(self.d_model)
+        x = self.pos_encoding(x)
+        x = self.transformer(x, mask=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,
+    ) -> 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}
+
+        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] + [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))
+                if target_len > self.model.pos_encoding.pe.size(1) - 1:
+                    target_len = self.model.pos_encoding.pe.size(1) - 1
+
+                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)
+                    if not torch.isfinite(probs).all() or probs.sum() == 0:
+                        break
+                    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[1:]])
+
+        return results

model.safetensors

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

requirements.txt

@@ -0,0 +1,4 @@
+torch>=2.12
+numpy>=1.26
+safetensors>=0.4
+gradio>=5.0