259 lines
9.8 KiB
Markdown
259 lines
9.8 KiB
Markdown
---
|
|
tags:
|
|
- text-generation
|
|
- transformers
|
|
- safetensors
|
|
- gguf
|
|
- text-generation-inference
|
|
- mixture-of-experts
|
|
- tipo
|
|
datasets:
|
|
- KBlueLeaf/danbooru2023-metadata-database
|
|
- KBlueLeaf/danbooru-8M-qwen3.5-caption
|
|
- CaptionEmporium/coyo-hd-11m-llavanext
|
|
- CaptionEmporium/laion-coco-13m-molmo-d-7b
|
|
- pixparse/cc12m-wds
|
|
arxiv: 2411.08127
|
|
language:
|
|
- en
|
|
pipeline_tag: text-generation
|
|
library_name: transformers
|
|
---
|
|
|
|
# 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](https://github.com/KohakuBlueleaf/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
|
|
|
|
<!-- PLACEHOLDER: confirm extension support for v2 before publishing -->
|
|
|
|
```python
|
|
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:
|
|
|
|
```bibtex
|
|
@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:
|
|
|
|
```bibtex
|
|
@software{kohakuwullm,
|
|
title={KohakUwULLM: an extensible decoder-only LLM training framework},
|
|
author={Yeh, Shih-Ying},
|
|
url={https://github.com/KohakuBlueleaf/KohakUwULLM},
|
|
year={2026}
|
|
}
|
|
``` |