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
- Why not just turn up CFG?
- Usage
- Recommended settings
- Temperature and min_p
- Rendering characters the image model doesn't know
- Repository contents
- Model details
- Training
- 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:
| you type | what it adds |
|---|---|
1girl, solo, school_uniform |
ribbon, upper_body, red_hair, serafuku, female_focus |
scenery, forest, sunlight |
water, mountain, lake, pond, house, landscape, cloud, cloudy_sky, sky, outdoors, nature, no_humans |
1girl, kimono, cherry_blossoms |
japanese_clothes, pink_kimono, wide_sleeves, petals, twintails, bow, ribbon, hair_ribbon, open_mouth, absurdres |
| nothing at all | 1girl, purple_hair, red_eyes, ponytail, maid_headdress, detached_collar, japanese_clothes, alternate_costume, upper_body, looking_to_the_side |
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.5and row 2temp=0use 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
| name | meaning |
|---|---|
temp |
sampling temperature, applied within the min_p support. 0 is greedy and deterministic |
min_p |
the support: tags below min_p * max(p) at temp=1 are never sampled, whatever the temperature. Default 0.01 |
types_mask |
which gelbooru categories may be emitted. Defaults to NO_ARTIST; note rank() and complete() default to ALL_TYPES instead |
ban |
tags never sampled. Suppressed during sampling, so the set completes around them |
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.generalonly: 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:
| node | what it does |
|---|---|
| S2SRec2 Loader | picks the checkpoint and vocabulary, caches them, outputs an S2SREC2_MODEL |
| S2SRec2 Tag Complete | the main one. Wire prompt into a CLIPTextEncode |
| S2SRec2 Tag Suggest | top-k next tags plus the completeness probability |
| S2SRec2 Info | what checkpoint is currently loaded |
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:
| temp | 0 | 0.5 | 0.75 | 1.0 | 1.25 | 1.5 | 1.75 | 2.0 | 2.25 |
|---|---|---|---|---|---|---|---|---|---|
| distinct sets per 64 | 1 | 23 | 43 | 56 | 60 | 64 | 64 | 64 | 64 |
| mean tags generated | 2.0 | 2.7 | 3.7 | 5.7 | 7.6 | 8.7 | 11.7 | 12.3 | 11.9 |
| mean pairwise Jaccard | 1.00 | 0.57 | 0.33 | 0.17 | 0.13 | 0.10 | 0.09 | 0.09 | 0.08 |
| distinct tags explored | 2 | 30 | 61 | 112 | 152 | 178 | 260 | 273 | 273 |
(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):
| temp 1.5 | 2 | 3 | 4 | 6 | |
|---|---|---|---|---|---|
min_p=0 (pure tempering) |
50 | 48 | 48 | 51 | 51 |
min_p=0.01 (default) |
27 | 29 | 39 | 51 | 53 |
min_p=0.02 |
22 | 22 | 33 | 37 | 47 |
min_p=0.04 |
18 | 16 | 20 | 24 | 23 |
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
copyrightfor 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, sotypes=generalgives 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
| Encoder | learned tag embeddings, then 2x ISAB (32 inducing points), giving contextual set Z |
| Retrieval head | PMA_1(Z) to a query, then a decoder over the emittable vocabulary |
| Completeness head | PMA_1(Z) to an MLP, one logit |
d_model / heads |
256 / 8 |
| Input vocabulary | 638,752 tags |
| Emittable vocabulary | 232,145 tags |
| Training corpus | 13,496,920 booru posts |
| Checkpoint | epoch 29 of run 20260725-035152 |
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
| category | tags | share |
|---|---|---|
| artist | 297,018 | 46.5% |
| character | 165,459 | 25.9% |
| general | 141,183 | 22.1% |
| copyright | 34,791 | 5.4% |
| meta | 301 | 0.0% |
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.
| category | in vocabulary | emittable | % | input-only |
|---|---|---|---|---|
| general | 141,183 | 141,183 | 100% | 0 |
| meta | 301 | 301 | 100% | 0 |
| character | 165,459 | 35,851 | 21.7% | 129,608 |
| artist | 297,018 | 49,680 | 16.7% | 247,338 |
| copyright | 34,791 | 5,130 | 14.7% | 29,661 |
| total | 638,752 | 232,145 | 36.3% | 406,607 |
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:
artistnames 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.copyrightis mostly the same, with franchise names that a base model may have partial exposure to.characterbreaks 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.generalbreaks completely, and gets no floor at all. These are English words. The text encoder parsesbound_to_treecompositionally 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:
| row | completeness label | retrieval target |
|---|---|---|
| partial basket | 0 | the dropped tags |
| complete set | 1 | none |
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:
| allowed types | mean tags generated | median |
|---|---|---|
general |
4.3 | 3 |
general,meta |
4.4 | 3 |
general,character,copyright |
9.3 | 7 |
general,character,copyright,meta |
9.4 | 8 |
| all five | 11.5 | 9 |
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.
| reference code | here | why | |
|---|---|---|---|
| Retrieval loss | single-target cross-entropy, leave-one-out | multi-positive over the dropped set | autocomplete wants a distribution over all plausible completions |
| Masking | leave-one-out, target is the highest tf-idf tag | random 20-60% drop | users type few tags, so baskets are sparse rather than nearly complete |
| Embeddings | frozen BERT-768 with a linear reduction | trainable learned table | tags are not natural language; BERT on 1girl or aigis_(persona) is weak |
| Padding | not masked, pad embedding joins attention | masked everywhere | correctness; the unmasked behaviour looks like a bug |
| alpha | 0.8 in code | 0.6 | the value their paper states |
| Pooling | PMA(k=2), two SAB layers, flatten, linear |
PMA(k=1) per head, one forward pass |
one seed sufficed, and two pools give each head its own view |
| Type conditioning | none | mandatory | callers suppress whole categories at serving time |
| Decoder bias | default init | vocabulary log-odds | starts calibrated to base rates |
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.



