TIPOv2-1B-A200M

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

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:

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.

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.

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:

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

Usage

<!-- PLACEHOLDER: confirm extension support for v2 before publishing -->

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

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