s2srec2-booru-tags

S2SRec2: booru tag completion

Give it a few booru tags, get back a complete tag set. It decides for itself when the set is finished, so you never pick a length.

Contents

• Overview
• Examples
  • Temperature
  • Seed tags
  • Allowed tag categories
• Why not just turn up CFG?
• Usage
  • Python
  • ComfyUI
  • Server
• Recommended settings
• Temperature and min_p
• Rendering characters the image model doesn't know
• Repository contents
• Model details
  • Vocabulary
  • Input and output vocabularies are different sizes
    • How the floors were chosen
• Training
  • Objective
  • Supervision
  • Type conditioning
  • Hyperparameters
  • Deviations from the reference implementation
• Credits
• License
• Provenance and scope

Overview

It exists to write prompts for image generators. Type 1girl, solo, school_uniform and it returns the tags a real post with those tags would have carried, which is a far better prompt than the three you typed. It also works as a live autocomplete: rank the most likely next tags given what is already in the box.

Real draws, one per starting set, at the recommended temp=1.5:

Neither length was specified. The model decides when a set is finished, and a landscape takes more tags to complete than a portrait. The scene row also produced no_humans unprompted. Every call returns a different set.

Two things it is not. It is not a text model, so it does not read or write sentences, only booru tags. And it has no safety filtering of its own; the vocabulary is whatever the corpus contained.

Examples

Every image below was produced end to end by this model driving an anime diffusion model. The renderer is held fixed throughout (1024x1024, 40 steps, CFG 4.5, Euler a), so the only thing changing is the tag model. Captions show the tags you typed in grey and what the model added in red.

Temperature

Seven temperatures from 0 to 2.0, four draws each, same seed tags throughout. The top row is greedy sampling, which returns the identical prompt every call, so all four images share one description and differ only by diffusion noise. By the bottom row the model is writing a different scene each time.

Seed tags

What you type is what it completes from. 1girl, solo, school_uniform gives school scenes, scenery, forest, sunlight gives landscapes with no people in them at all, and an empty seed gives the unconditioned base rate. Temperature is fixed at 1.5.

Allowed tag categories

general only against general plus meta, same seed and temperature. Narrowing what the model may emit changes where it decides the set is finished, not just which tags survive.

examples/INDEX.txt lists all 52 prompts. examples/manifest.json carries the full settings behind every image.

Some draws were regenerated to keep the published set SFW. That is a hosting constraint, not a description of the model, whose output is unfiltered and includes NSFW content.

Why not just turn up CFG?

Image models already have a diversity knob. Turning CFG down makes output vary more, so it is worth being precise about what a prompt model adds.

Both knobs are temperature-like. Ho and Salimans frame guidance as trading mode coverage for fidelity "in the same spirit as low temperature sampling". But the two sharpen different distributions. CFG sharpens p(image | prompt) with the prompt held fixed. Tag temperature flattens p(tag set | seed), which changes the prompt itself.

Guidance extrapolates along (eps_cond - eps_uncond), a direction the conditioning already defines. No CFG value can introduce a concept the prompt does not contain. CFG also has no notion of prompt length, which the stop head varies by 5x across its range.

A controlled comparison. One diffusion seed is pinned across all nine renders, so within a row the only variable is the axis under study.

• Row 1, CFG axis. The tag set is frozen and cfg moves from 1 to 16.
• Row 2, temperature axis. cfg is frozen at 4.5 and tag-model temperature moves from 0 to 2.0.
• Control. Row 1 cfg=4.5 and row 2 temp=0 use the same prompt, cfg and seed. They render byte-identically (verified by SHA256), which is what makes the rows comparable.

Row 1 never changes the subject. Every cell is the same girl in the same serafuku, and cfg only changes how hard that description is enforced. At cfg=1 the image collapses into incoherent mush, which is not creativity, just failure to follow the prompt. From 4.5 through 16 the same scene grows harder, more saturated, then flat and artifacted.

Row 2 changes the subject in every cell: a classroom desk scene, an orange-haired character, a green-haired one outdoors with a drink. Same fidelity throughout. What moved is what was asked for.

CFG modulates conviction. This modulates content. They compose, and neither replaces the other.

Usage

Everything needed to run the model is in this repository. Three ways in: call it from Python, drop the ComfyUI node pack into your workflow, or run the reference HTTP server.

Python

Needs only torch and numpy.

from s2srec2._imports import Vocab
from s2srec2.infer import load_model, rank
from s2srec2.generate import TagGenerator
from s2srec2.typemask import NO_ARTIST

vocab = Vocab.load("vocab_s2srec2.json")
model, cfg = load_model("tag_s2srec2.pt", vocab, "cpu")
gen = TagGenerator(model, vocab, "cpu")

gen.generate(["1girl", "solo", "school_uniform"])
# ['blush', 'thighhighs', 'serafuku', 'highres', 'skirt']      varies per call

gen.generate(["1girl", "solo", "school_uniform"], temp=0.0)
# ['female_focus', 'serafuku']                                  greedy, deterministic

gen.generate([], ban=["monochrome"], types_mask=NO_ARTIST)      # no seed, one banned tag

TagGenerator.generate() is the main entry point. It samples with temp and min_p, stops on the completeness head, suppresses junk tags, disallowed categories and your bans during sampling, and returns only the new tags. Your seed stays yours to keep and append to. Pass a torch.Generator for reproducible draws. generate_set(model, vocab, seed, ...) is a one-shot wrapper that caches per-vocabulary tensors for you.

Also available: infer.rank() returns the top-k next tags plus the completeness probability, which is what an autocomplete box wants. infer.complete() grows a set greedily, taking the argmax every step, so it returns the same answer every call.

Parameters

types_mask names what stays allowed, not what is removed. A mask that forgets a category silently generates without it. The model's forward pass itself refuses to run without an explicit mask; it is the Python entry points that fill in the defaults above. s2srec2.typemask provides ALL_TYPES, NO_ARTIST and helpers. Common choices:

• NO_ARTIST: everything except artists. Artist choice usually belongs to the user, not the model.
• general only: no character, copyright or meta tags. Use this when output must not contain IP.

ComfyUI

comfyui-s2srec2/ in this repository is a ready to use node pack. Copy it into ComfyUI/custom_nodes/, put tag_s2srec2.pt and vocab_s2srec2.json into ComfyUI/models/s2srec2/, and restart ComfyUI.

ComfyUI/
  custom_nodes/comfyui-s2srec2/     <- this folder
  models/s2srec2/
    tag_s2srec2.pt
    vocab_s2srec2.json

Nothing else to install. The model package is vendored inside the node and needs only torch and numpy, which ComfyUI already has.

Four nodes appear under the s2srec2 category:

Tag Complete takes the same knobs described above (temperature, min_p, per-category switches, ban) plus a seed with control_after_generate, so each queue produces a fresh set the way a KSampler does. It outputs prompt, added_tags and unrecognised.

It reads a real ComfyUI prompt rather than requiring a clean tag list. Spaces, mixed case, weights like (masterpiece:1.2), emphasis brackets, <lora:...>, BREAK and embedding: are all handled, commas are optional, and a natural language sentence has its tags picked out of it. Your own text is passed through byte for byte and the generated tags are appended after it; anything that could not be resolved to a tag is reported, per segment, on the unrecognised output rather than being silently dropped. See comfyui-s2srec2/README.md for the full parsing rules, and comfyui-s2srec2/demo_parsing.py for a runnable demonstration across every input format.

Server

serve_reference.py is a Flask server exposing /api/random_prompt, /api/suggest and /api/health.

pip install -r requirements.txt
python serve_reference.py
curl "http://127.0.0.1:8000/api/random_prompt?seed=1girl,solo,school_uniform&n=2&temp=1.5"

{"prompts": [["serafuku","female_focus","long_hair","blush"], ["skirt","brown_eyes"]],
 "seed": ["1girl","solo","school_uniform"], "types": "no:artist", "unknown": []}

The seed is not echoed back in prompts. unknown reports seed tags outside the vocabulary, which are ignored rather than treated as errors.

The server uses ONNX Runtime for its forward pass. On first start it exports tag_s2srec2.onnx (plus a .onnx.data of roughly 870 MB) next to the checkpoint, which takes a while and needs the disk space. That file is not shipped, so it always matches the checkpoint you are running. Delete it after swapping models to force a re-export. The direct Python path above does not use ONNX at all.

Recommended settings

Use temp=1.5 with the default min_p=0.01. Measured over 64 draws per point, NO_ARTIST:

(Anchored seed 1girl, solo, school_uniform. An empty seed has the same shape, shifted longer. Raw data in training/temp_curve.json.)

1.5 is the lowest temperature at which the model is fully non-degenerate, 64 distinct sets out of 64, while length stays controlled at a median of 7 tags.

Below 1.0 it collapses. At temp=0 all 64 draws return the same set from a two-tag vocabulary. temp=1.0 looks varied (56 of 64) but explores only 112 distinct tags against 178 at 1.5. Above 1.5 exploration keeps widening while anchored length plateaus around 12 tags; how far the range safely extends is min_p's decision, covered in the next section.

Temperature and min_p

Sampling is support-pinned: min_p truncates the temp=1 distribution before temperature is applied, so the set of tags that may be sampled — the rules of what a set may contain — never changes with temperature. This matters because rarity and incompatibility both show up as small probabilities, and tempering first would flatten that distinction away, letting high temperature admit tags the conditional had ruled out (short_hair beside medium_hair, four incompatible skirts). With the support pinned, temperature only redistributes preference among tags the model already considers plausible, so a high-temperature set is an unusual set, not an incoherent one.

The two knobs therefore separate cleanly. temp decides how adventurously a draw roams within the support. min_p decides what the support is — and with it, where generation can still run away into very long sets. The stop head is a completeness predicate trained on real posts: give sampling enough of the tail (low min_p) at high temperature and the growing set stops resembling any finished post, so the head holds out until sheer size makes it read as complete, near the top of the training corpus's max_tags=64 window.

Measured unseeded, the worst case (mean tags generated, training/runaway_onset.json):

At the default the wild zone starts around temp=3 and reaches the max_steps=60 cap by 4. At min_p=0.04 it never arrives (measured to temp=6); at min_p=0 it is already there by 1.5. Anchored seeds are far tamer — real seed tags hold the set on the manifold, and the school-seed curve above plateaus around 12 tags. Pick min_p for how much tail you want reachable and where the wild end of your temperature range should sit; 0 recovers pure tempering, and the pre-fix sampler's behaviour is recorded in training/stop_probe.json.

Rendering characters the image model doesn't know

An image model can only draw a character it memorised. Ask for one it never saw and you get a generic figure, because the name is an unfamiliar token. The tag model knows what that character looks like, not as an image but as a set of co-occurring tags: whoever tagged those posts recorded the hair colour, the eyes, the outfit. Seeding the tag model with an unknown character name and feeding the expansion to the renderer gets you something much closer than the bare name does.

This is what the asymmetric vocabulary buys. The 406,607 input-only tags are almost entirely proper nouns below the emission floors, so they can be conditioned on but never emitted. The model takes a rare character name in and answers in general tags, which are fully emittable and which the renderer reads compositionally.

Caveats:

• Suppress copyright for this. Expansions often surface the series tag, which is emittable even when the character is not. That pulls the render toward the series house style and away from the specific character, which can produce a series-typical outfit the character never wears. The character information lives in the attribute tags, so types=general gives the cleanest transfer.
• It needs the character to be well represented in the corpus. A tag with a handful of posts has a conditional too thin to describe anything.
• Only describable attributes transfer: hair, eyes, clothing, accessories. Face shape, art style and anything the tag vocabulary has no word for do not.
• Fidelity is "closer", not "correct". It is a description, not a likeness.
• Raising temperature weakens it. The seed stays in the prompt, but competing sampled tags can overwhelm it.

How well this works varies enough between characters that the tag output is a better guide than any single rendered example.

Repository contents

tag_s2srec2.pt              checkpoint (epoch 29, run 20260725-035152)
vocab_s2srec2.json          tag names, categories, counts. Required to load the model
calibration.json            display calibration, all-types fit
calibration.no_artist.json  display calibration, artist-suppressed fit
serve_reference.py          reference inference server
requirements.txt
build_dataset.py            corpus -> vocab/pack training bundle, one command (see TRAINING.md)
build_output_head.py        emittable-tag list from the per-category count floors
combine_dataset.py          de-duplicating corpus merge
comfyui-s2srec2/            ComfyUI node pack, self contained
s2srec2/                    the model package: inference and training
cvae/                       vocabulary, tag filtering and packing helpers
TRAINING.md                 how to train, and what data is not included
training/                   run hyperparameters, training curves, temperature / type-mask /
                            stop-behaviour measurements
examples/                   61 rendered examples (52 sweep + 9 CFG comparison), contact sheets,
                            per-image settings

s2srec2/ and cvae/ must stay siblings; s2srec2/_imports.py puts cvae/ on the path. serve_reference.py resolves the checkpoint, vocabulary and calibrations from its own directory, so the repository runs as-is after git lfs pull.

Model details

Both PMA pools read the same Z, so the completeness head sees the type condition as well as the tags. Training curves are in training/.

Vocabulary

That is a share of the vocabulary, not of posts. Artist names are numerous but each is rare. There are five categories; gelbooru's deprecated was resolved away at build time, with 19 head tags re-typed by hand (about 674k occurrences corrected) and the rest folded into general.

Input and output vocabularies are different sizes

The model accepts 638,752 tags and can emit 232,145.

A rare tag is valuable to condition on and close to worthless to emit. The softmax is frequency-dominated by construction, since the decoder bias is initialised from vocabulary log-odds, and a count-10 tag was measured reachable in a top-50 ranking only about 9% of the time. So the embedding covers everything down to min_count=3 while the decoder covers only tags worth suggesting.

How the floors were chosen

The question behind a floor is not "is this tag rare?" but "does the renderer do anything when it sees this tag?" A tag the renderer ignores is wasted softmax mass, and worse, it is a tag the sampler can spend a slot on instead of one that works.

You cannot measure that directly, because an image model's training set is not published. What you can use is the corpus count as a one-sided bound. If the renderer was finetuned on booru data drawn from the same tag universe, then for any given tag its exposure is at most our count. So if tags at count N produce no visible effect, nothing below N can be better learned, and excluding count <= N is safe. That argument only runs downward. A high count does not promise a tag works, which is why this is used to delete tags and never to justify keeping them.

The bound only holds where the renderer learned the concept from booru data in the first place, and that is what forces the floors to be per category rather than one global number:

• artist names are novel to a general-purpose base model, so the finetune is the only source and the count bounds it cleanly. This is also the largest slice of the vocabulary, so it is where a floor buys the most.
• copyright is mostly the same, with franchise names that a base model may have partial exposure to.
• character breaks for the famous head. A base model may already know a widely drawn character from pretraining, so a low booru count does not mean low exposure.
• general breaks completely, and gets no floor at all. These are English words. The text encoder parses bound_to_tree compositionally whether or not that exact tag was ever a training label, so the count is not evidence of anything.

Judging whether a tag registers is done by eye, on rendered output, because there is no classifier that can do it in the long tail where the cutoff sits. The signal is consistency rather than correctness: render the tag many times, and one the renderer knows produces the same distinctive subject or style across samples, while one it does not know collapses to the model's default look. That collapse is the reliable tell, and it means a tag can be judged even when you do not personally recognise the character or artist behind it. Working down from a few hundred, per category, until the effect stops being visible is where the floors come from: artist >= 50, character >= 50, copyright >= 150.

Eyeballing alone would be unreliable, because response varies a great deal between tags at the same count, and judging a handful that happen to be unusually strong would put the floor too low and delete working tags. A numeric per-tag response was measured on the renderer and regressed against log corpus count within each category, and a tag was only put in front of a human if its measured response sat close to the fitted line for its count. Outliers were swapped out. That measurement selected which tags to judge; it did not set the cutoff.

This is an approximation. The floor is a count threshold applied to a population, while the thing that matters is per-tag response, and count is only a proxy for it. Excluding on the measured response directly, tag by tag, would be better than converting it into a count cutoff, but that requires measuring the whole vocabulary rather than a sample.

The numbers are also specific to one renderer. A different image model will have its floors somewhere else, and moving them means repeating the judgement against that model.

Dropping 63.7% of the vocabulary from the output side still leaves 98.8% of all corpus tag occurrences covered, because everything removed is rare by definition. The narrowing also pays for itself in training, where the decoder matrix multiply and the full softmax dominate step cost and both scale with output width.

Training

Objective

L = alpha * L_retrieval + (1 - alpha) * L_completeness        alpha = 0.6

L_retrieval is a masked softmax cross-entropy over the whole set of dropped tags, divided by the number of them. The softmax denominator is restricted to the categories the type mask allows, so the model learns P(tag | basket, mask) and is never scored on tags the condition forbids. L_completeness is binary cross-entropy on one logit.

Supervision

Each record's tag set is the target. Per training example a type mask is drawn first and applied to the whole record, then a random fraction of tags is dropped to form a partial basket, with the dropped tags as the retrieval target. Each example expands into two encoder rows:

The drop fraction is 20% to 60%, and a basket can be as small as one tag, so the model learns completion from sparse input rather than only from nearly-finished sets.

Applying the mask before the split is what makes "complete" mean complete under this condition. The partial basket, the retrieval target and the complete row all come from the same type-filtered set, so when artists are switched off the model learns what a finished set looks like without one.

Type conditioning

Every tag carries a gelbooru category, and callers routinely suppress whole categories: artists because that choice belongs to the user, characters and copyrights when output must not contain IP. The model takes the allowed-category mask as an input so it can generate correctly under those conditions.

The mask is embedded and prepended to the set as a token before the ISAB stack, so the encoder sees it and both pooling heads inherit it. Training samples uniformly over all 31 non-empty allowed subsets of the five categories, so no condition is better supported than another.

The flag is derived from the intervention, not from the data. A category is dropped at random and the bit is set accordingly. A post that happens to have no artist tag is still labelled artist-allowed, because "allowed" means may appear, not does appear. Conditioning on natural absence would bias generation toward the kinds of posts that lack artists, such as scans and screencaps, whenever artists were switched off.

Each tag token also carries its own category through a small embedding summed onto the tag vector. The category is known metadata, so making the model infer it from co-occurrence would spend capacity for at best the same result.

The effect is measurable. Same seed (1girl, solo, school_uniform), same temperature (1.5), same default min_p, only the mask moving, 64 draws per condition:

Opening up character and copyright roughly doubles the sets, and artist extends them again: the stop head holds out for the kinds of tags the wider condition allows. Meta barely moves the needle. Raw data in training/types_curve.json.

Hyperparameters

d_model 256, 8 heads, 32 inducing points, 2 ISAB layers, PMA k=1 per head, dropout 0.1, batch 512, lr 1e-4, 30 epochs, bf16 with tf32, min_tags 4, max_tags 64, learned embeddings, decoder bias initialised from vocabulary log-odds. Full configuration in training/train_config.json.

Deviations from the reference implementation

The authors' released code was inspected and not followed where it contradicts their published paper; the paper's described method was implemented instead. The full record with rationale is in s2srec2/DESIGN_AND_DEVIATIONS.md.

Sequence length is dynamic rather than fixed at 45. Evaluation adds a calibration suite, measuring Spearman correlation against the corpus's empirical P(tag|query) plus attraction and hallucination@k, on the principle that a prediction is only wrong if it exceeds the true conditional.

Credits

This implements the method described in S2SRec2, applied to a different domain. The paper's task is grocery basket completion: given a partial set of ingredients, predict the complementary set. The reformulation here is the same with ingredients replaced by booru tags. The multitask framing, the Set Transformer backbone, the sequential add-until-stop inference and alpha = 0.6 all come from that paper.

S2SRec2: Set-to-Set Recommendation for Basket Completion with Recipe Yanan Cao, Omid Memarrast, Shiqin Cai, Sinduja Subramaniam, Evren Korpeoglu, Kannan Achan arXiv:2507.09101

The encoder is built from Set Transformer blocks (MAB, SAB, ISAB, PMA), which give the model permutation invariance, the property that makes it correct for sets rather than sequences.

Set Transformer: A Framework for Attention-based Permutation-Invariant Neural Networks Juho Lee, Yoonho Lee, Jungtaek Kim, Adam R. Kosiorek, Seungjin Choi, Yee Whye Teh ICML 2019, arXiv:1810.00825

Type conditioning, the per-token category embedding, learned tag embeddings, fraction-based masking and the calibration suite are not from either paper. Neither has any notion of tag type, since their vocabulary is one kind of thing.

@article{cao2025s2srec2,
  title  = {S2SRec2: Set-to-Set Recommendation for Basket Completion with Recipe},
  author = {Cao, Yanan and Memarrast, Omid and Cai, Shiqin and Subramaniam, Sinduja
            and Korpeoglu, Evren and Achan, Kannan},
  journal = {arXiv preprint arXiv:2507.09101},
  year   = {2025}
}

@inproceedings{lee2019set,
  title     = {Set Transformer: A Framework for Attention-based Permutation-Invariant
               Neural Networks},
  author    = {Lee, Juho and Lee, Yoonho and Kim, Jungtaek and Kosiorek, Adam R.
               and Choi, Seungjin and Teh, Yee Whye},
  booktitle = {Proceedings of the 36th International Conference on Machine Learning (ICML)},
  pages     = {3744--3753},
  year      = {2019}
}

License

CC0 1.0 Universal, a public domain dedication. See LICENSE.

No copyright is asserted over this release. A permissive licence such as MIT would be the wrong instrument, since it grants rights from a copyright holder, and the interesting part of this model is not ours to hold. What it learned came from a scraped Gelbooru corpus and from the people who made and tagged those images. CC0 waives whatever rights might attach to the code and weights and gets out of the way.

This is not a warranty that any particular downstream use of the output is clear. CC0 explicitly disclaims responsibility for clearing third-party rights. That assessment is yours.

Provenance and scope

Trained on 13.5M booru posts. The vocabulary reflects that corpus and its biases: 46.5% of entries are artist names, and a substantial share of the emittable vocabulary is NSFW.

The model has no safety filtering. types and ban are the only levers and neither is a content filter. Rating tags such as safe are not in the vocabulary and do not condition generation at all, though passing one through to a downstream image model that does know it will steer the render.

Character and copyright tags will reproduce properties from the training corpus. Suppress those categories through types if that matters for your use.