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

TIPOv2-1B-A200M is a 1B-A200M sparse model. 991M total parameters, ~200M active per token, plus a ~50M embedding table. The second generation of TIPO, rebuilt from the dataset up on the KohakUwU MoE architecture.

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 narrow the range indicate by a brief prompt to a more specifici description which match the original prompt, allow diffusion model to have more information to work while persist 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. It is not a fine-tune of TIPO-500M. The dataset, the captioner and the architecture are all different.

What makes it v2

1. Fully upgraded dataset

	v1 (TIPO-500M)	v2 (TIPOv2-1B-A200M)
sources	GBC10M, Danbooru, CoyoHD-11M	latest Danbooru, Nozomi, CC12M, CoyoHD-11M, LAION-COCO-13M
breadth	3 sources	5 sources, both anime-domain and general-photography

The v1 mix leaned heavily on one general-caption source. v2 adds Nozomi and LAION-COCO-13M alongside a refreshed Danbooru, which broadens both the tag vocabulary and the image domains the model has seen. Danbooru is weighted x3 and the dedicated tagger view x2, so booru-style tag structure stays dominant while the general sources supply natural-language variety.

2. Better natural-language captions

Every natural-language caption in v2 is regenerated with Qwen3.5-2B. In v1 the caption quality varied by source, because each dataset shipped whatever captions its authors produced. Regenerating them under a single captioner means caption style is constant across sources, and the only thing that varies between coyo11m and laion_coco is the image distribution, not the writing. That makes the source weighting a choice about visual domain rather than an accidental choice about prose quality.

3. Fully upgraded architecture: KohakUwU MoE

v1 was a 500M dense LLaMA-like arch. v2 uses the KohakUwU MoE architecture, 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. TIPOv2 is one model trained with it.

The configuration used here:

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. So v2 carries roughly 2x v1's parameters while activating fewer of them per token than v1's dense 500M. Context is 4096, up from v1's 1024.

Training recipe

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

steps	150,000
tokens per step	262,144 (16384 x 16 microbatches)
context	2048 packed
parallelism	4-stage pipeline, 1F1B schedule
parameter dtype	full fp16 (with dynamic loss scaling)
autocast	fp16
MXFP8	q/k/v/o projections and MLP up/down (w_in/w_out), including the shared expert. 111 modules.
routed experts	fused MXFP8 expert path
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	2% of run (3000 steps)
grad clip	1.0
aux loss / router z-loss	0.0 / 0.0, since balancing is aux-loss-free

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 2048 context) a padded batch would be ~80% 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; the run reports zero overflows at scale 65536.
• 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.
• MXFP8 on the dense projections. Block-scaled fp8 (E4M3 with a shared power-of-two scale per 32 elements) on q/k/v/o and up/down. The routed experts use a fused MXFP8 path whose epilogues never materialize the (tokens x top_k, hidden) intermediates.

Tokenizer

The tokenizer is 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 65535	1519	<|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. 1519 slots are still free in this release.

Prompt format

quality: masterpiece
rating: general
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.

Usage

from transformers import AutoTokenizer, AutoModelForCausalLM

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

prompt = (
    "quality: masterpiece\n"
    "rating: general\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
hf/model.safetensors	1.98 GB	transformers, fp16
gguf/TIPOv2-1B-A200M-f16.gguf	2.02 GB	llama.cpp, fp16
gguf/TIPOv2-1B-A200M-Q8_0.gguf	1.07 GB	llama.cpp, 8-bit

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