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:

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:

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

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.

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.

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:

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.

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:

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

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}
}