TIPO-v2.1-1B-A200M: Next generation of T2I prompt optimization model.

TIPO-v2.1-1B-A200M is a 1B-A200M sparse model. 991M total parameters, ~200M active per token, plus a ~50M embedding table. It is a full retrain of TIPOv2 on a corrected dataset, for longer, at 4096 context, in pure fp16.

Introduction

TIPO is a framework for improving Text-to-Image generation by text presampling: a small language model expands a short user prompt into a detailed one before the diffusion model ever sees it. Pre-sampling a smaller distribution by narrowing the range indicated by a brief prompt to a more specific description which matches the original prompt allows the diffusion model to have more information to work with, while preserving overall diversity and fidelity. Instead of asking the user to write 200 tokens of booru tags and natural language, TIPO samples that expansion from a distribution learned over real caption data.

This is v2.1. Same architecture and same data sources as v2, retrained from scratch after three defects were found in how v2's training examples were built.

What's new in v2.1

1. Metadata that was silently ignored now works

v2's renderer emitted three metadata fields in forms the model could not learn from. Each is fixed at the source in v2.1, so rating, quality and person-count tags actually condition generation.

rating was out of vocabulary. v2 passed Danbooru's raw values straight through, but the TIPO prompt vocabulary (and everything built on it, including KGen) uses different words. A caller asking for rating: safe was sending a token the model had never been trained on. v2.1 converts:

Danbooru	TIPO
general	safe
sensitive	sensitive
questionable	nsfw
explicit	nsfw, explicit

quality was badly miscalibrated. v2 bucketed quality by fixed thresholds on a post's raw score. Danbooru scores inflate over time and recent posts have not yet accumulated votes -- median score runs 21 (2022), 16 (2023), 6 (2024), 5 (2025) -- so fixed thresholds label almost everything recent as bad, and the problem compounds as the corpus grows. The realised v2 distribution was 51% low quality and 0.7% worst quality, plus a below average quality bucket that is not in the TIPO vocabulary at all.

v2.1 ranks each post against its own year, by percentile over both score and fav_count, so "masterpiece" means top-5% for the year it was posted:

label	target	realised
worst quality	10%	9.6%
low quality	15%	14.5%
normal quality	25%	25.6%
good quality	25%	25.4%
great quality	10%	9.9%
best quality	10%	9.9%
masterpiece	5%	5.1%

Person-count tags could be dropped. 1girl, 2boys, multiple girls and the rest are the strongest structural signal in a booru prompt. v2 treated them as ordinary general tags: shuffled with everything else, split arbitrarily between the prompt and the target, and discarded entirely when the renderer chose to drop input tags. v2.1 gives them their own category and always places them at the head of the tag list, including in the case where all other metadata and all seed tags are dropped:

target: <|very_short|>
tag: 2girls, multiple girls

2. Year is now part of the quality signal

Because quality is normalised per year, the era itself became information the model could no longer infer. v2.1 emits it alongside quality:

quality: masterpiece, newest

newest (2024+), recent (2020-2024), mid (2018-2020), early (2015-2018), old (before 2015).

3. Longer training, longer context

	v2	v2.1
steps	150,000	200,000
context	2048 packed	4096 packed
tokens seen	~39B	52.4B

4. Pure fp16 instead of MXFP8

v2 ran block-scaled fp8 (MXFP8) on 111 dense projection modules plus a fused MXFP8 expert path. v2.1 runs fp16 end to end with no quantised matmuls at all: on this model the fused 16-bit expert path is faster end to end than the MXFP8 one and carries no quantisation error, so the fp8 path was buying nothing. Loss scaling covers fp16's narrower exponent range; the run reports zero overflows across all 200,000 steps.

Architecture: KohakUwU MoE

Unchanged from v2. A DeepSeekMoE-style sparse decoder from KohakUwULLM.

KohakUwU is a series of projects for pretraining infrastructure. KohakUwULLM is the general-purpose LLM training project within that series, and it is where this architecture, the training framework and the kernels described below come from. None of it is part of TIPO, and none of it was built for TIPO. TIPO-v2.1 is one model trained with it.

Configuration
total params	990.8M
active params / token	193.1M (excludes the embedding lookup)
input embedding	50.3M (a gather, not a matmul, so not counted as active)
output head	50.3M
routed experts	854.1M total, 106.8M active at top-8
attention + shared expert + dense layer + router + norms	36.1M, all active
layers	16 (layer 0 dense, 15 MoE)
hidden size	768
attention	12 heads, 2 KV heads (GQA), head dim 64, QK-norm
routed experts	64, top-8 per token
shared experts	1 (always on)
expert hidden	384
dense MLP hidden	2048
router	sigmoid scoring, aux-loss-free bias balancing
position	RoPE, theta 100000, 4096 context
norm	RMSNorm, eps 1e-6
vocab	65536

Only 193M of 991M parameters do work on any given token. The 854M of routed experts contribute just 107M at top-8, and the 50M embedding is a lookup rather than a matmul.

Training recipe

Trained on 4x RTX 5090 (32 GB, sm_120) with KohakUwULLM, in 55 hours.

steps	200,000
tokens per step	262,144 (8192 x 32 microbatches)
total tokens	52.4B
context	4096 packed
parallelism	4-stage pipeline, 1F1B schedule
parameter dtype	full fp16 (with dynamic loss scaling)
autocast	fp16
quantised matmuls	none
optimizer	Muon on hidden matrices, AdamW on the rest
LR	5e-4 (muon_lr 2e-3, embed_lr 2e-3)
schedule	inverse-sqrt power (s0 2500, b -0.5), then cosine to 1%
warmup	2000 steps
grad clip	1.0
aux loss / router z-loss	0.0 / 0.0, since balancing is aux-loss-free
peak memory	11.02 GiB / rank

Final: loss 0.9902 (EMA), perplexity 2.703, 0 dead experts, 0 overflows.

Data mixture, unchanged from v2 -- Danbooru weighted x3 and the dedicated tagger view x2, so booru-style tag structure stays dominant while the general sources supply natural-language variety:

source	weight
Danbooru (metadata + Qwen3.5 captions)	3
Danbooru tagger view	2
CoyoHD-11M	1
LAION-COCO-13M	1
CC12M	1
Nozomi	1

Notes on the choices that are not obvious. All of these are KohakUwULLM facilities, not TIPO-specific work:

• Packed varlen, not padded. Every sequence is concatenated onto one flat token axis with cu_seqlens carrying document boundaries. For TIPO-shaped data (50 to 600 tokens against a 4096 context) a padded batch would be overwhelmingly padding.
• fp16 parameters, not bf16. fp16 carries 10 mantissa bits against bf16's 7. It needs loss scaling to keep its narrower exponent range in bounds, which the trainer supplies.
• Aux-loss-free balancing. Expert load is balanced by a selection-only bias updated outside the gradient, not by an auxiliary loss term. A router z-loss was measured at 1.59x end-to-end cost and left off.

Tokenizer

The DeepSeek-V4 tokenizer, pruned to 64000 ordinary BPE tokens, plus a 1536-slot block reserved for special tokens. Total vocabulary is 65536.

id range	count	contents
0 to 63999	64000	ordinary BPE tokens, kept in DeepSeek-V4 merge order
64000 to 64016	17	named specials: <|bos|>, <|eos|>, <|pad|>, <|unk|>, and the 13 TIPO control tokens
64017 to 64024	8	ChatML and tool-call specials, shared with other KohakUwULLM models and unused by TIPO
64025 to 65535	1511	<|reserved_N|> placeholders

Two reasons the layout looks like this:

• 65536 is a power of two. The output head is a GEMM whose N dimension is the vocabulary, and a power-of-two N keeps that GEMM tile-aligned. An awkward vocab size costs throughput on every token generated.
• The reserved block is deliberate headroom. Adding a control token later is an id assignment inside the existing embedding table, not a resize and re-embed. The ChatML block at 64017 is exactly that mechanism being used; 1511 slots are still free.

Prompt format

quality: masterpiece, newest
rating: safe
target: <|long|> <|tag_to_long|>
tag: 1girl, cherry blossoms, outdoors

Control tokens

Length targets, which set how long the generated result should be:

<|empty|> <|very_short|> <|short|> <|long|> <|very_long|>

Task selectors, which set what to generate from what:

token	meaning
<|tag_to_long|>	tags to long natural-language caption
<|long_to_tag|>	long caption to tags
<|short_to_tag|>	short caption to tags
<|short_to_long|>	short caption to long caption
<|tag_to_short_to_long|>	tags, then short, then long
<|short_to_tag_to_long|>	short, then tags, then long
<|short_to_long_to_tag|>	short, then long, then tags
<|gen_meta|>	also predict the metadata fields

Metadata lines the model understands, all optional: quality, rating, artist, characters, copyrights, meta, aspect ratio.

Vocabularies (changed in v2.1, see above):

• rating: safe, sensitive, nsfw, explicit
• quality: worst quality, low quality, normal quality, good quality, great quality, best quality, masterpiece, optionally followed by an era -- old, early, mid, recent, newest

Usage

from transformers import AutoTokenizer, AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "KBlueLeaf/TIPO-v2.1-1B-A200M", trust_remote_code=True, dtype="float16"
).cuda().eval()
tokenizer = AutoTokenizer.from_pretrained("KBlueLeaf/TIPO-v2.1-1B-A200M")

prompt = (
    "quality: masterpiece, newest\n"
    "rating: safe\n"
    "target: <|long|> <|tag_to_long|>\n"
    "tag: 1girl, cherry blossoms, outdoors\n"
)
ids = tokenizer(prompt, return_tensors="pt").input_ids.cuda()
out = model.generate(ids, max_new_tokens=256, temperature=1.0, min_p=0.1, do_sample=True)
print(tokenizer.decode(out[0], skip_special_tokens=False))

trust_remote_code=True is required, because the KohakUwU MoE architecture ships as modeling_kohaku.py beside the weights.

Files

file	size	use
model.safetensors	1.98 GB	transformers, fp16
TIPO-v2.1-1B-A200M-f16.gguf	2.02 GB	llama.cpp, fp16
gguf/TIPO-v2.1-1B-A200M-f16.gguf	2.02 GB	the same file, under gguf/
gguf/TIPO-v2.1-1B-A200M-Q8_0.gguf	1.07 GB	llama.cpp, 8-bit

The transformers files sit at the repository root, so from_pretrained needs no subfolder argument. The f16 GGUF is mirrored at the root for downloaders that cannot fetch from a subfolder.

The GGUF files carry no chat template: TIPO is a completion model and the prompt format above is the interface.

LICENSE

Released under Kohaku License 1.0.

Citation

TIPO:

@misc{yeh2024tipotextimagetext,
  title={TIPO: Text to Image with Text Presampling for Prompt Optimization},
  author={Yeh, Shih-Ying and Park, Sang-Hyun and Oh, Giyeong and Song, Min and Yu, Youngjae},
  year={2024},
  eprint={2411.08127},
  archivePrefix={arXiv}
}

The architecture, training framework and kernels:

@software{kohakuwullm,
  title={KohakUwULLM: an extensible decoder-only LLM training framework},
  author={Yeh, Shih-Ying},
  url={https://github.com/KohakuBlueleaf/KohakUwULLM},
  year={2026}
}