Unofficial ONNX export of pixai-labs/pixai-tagger-v1.0 in FP32, FP16 and INT8.
Runs with ONNX Runtime only.
Files
| File |
Precision |
Size |
model.onnx |
FP32 |
1.89 GB |
model_fp16.onnx |
FP16 (FP32 input/output) |
947 MB |
model_int8.onnx |
Dynamic INT8 (MatMul/Gemm, per-channel) |
523 MB |
selected_tags.csv |
tag_id,name,category,count |
preprocess.json |
Preprocessing steps |
- Input
input: float32 [batch, 3, 1008, 1008], RGB, NCHW.
- Output
prediction: float32 [batch, 30877], probabilities (sigmoid is part of the graph).
- Opset 17.
Test results
- Hardware: AMD Ryzen 5 9600X (6 cores / 12 threads, AVX-512 VNNI), 32 GB RAM, AMD Radeon RX 580X 4 GB.
- Software: Windows 11,
onnxruntime-directml 1.24, batch 1.
- Images: 6 anime illustrations. 5 of them are Danbooru posts uploaded in September 2026, after the model's
training cutoff, and are also scored against their Danbooru tags.
- Thresholds: the per-category values listed under Categories and thresholds.
- FP16 is meant for GPUs. On CPU it is slower than FP32, because the CPU has no native FP16 math.
- FP32 fits on a 4 GB GPU only barely (3.2 GB VRAM), and its speed varies between runs.
- INT8 speed depends on CPU support for VNNI instructions. On older CPUs it can be slower than FP32.
- Peak RAM on CPU includes session initialization; FP32 briefly reaches 9.4 GB while loading.
Agreement with FP32 (CPU)
| Model |
Device |
Max prob. diff |
Mean prob. diff |
Tag-set similarity (Jaccard) |
Tags lost / added (of 317) |
| FP32 |
DirectML |
6 × 10⁻⁶ |
2 × 10⁻⁸ |
1.000 |
0 / 0 |
| FP16 |
CPU |
0.002 |
7 × 10⁻⁷ |
1.000 |
0 / 0 |
| FP16 |
DirectML |
0.012 |
6 × 10⁻⁶ |
0.997 |
1 / 0 |
| INT8 |
CPU |
0.103 |
6 × 10⁻⁵ |
0.951 |
4 / 13 |
Against Danbooru tags (5 posts)
| Model |
General P / R / F1 |
Character P / R |
Copyright P / R |
Rating |
| FP32 |
0.50 / 0.73 / 0.59 |
6/6 · 6/8 |
6/7 · 6/6 |
5/5 |
| FP16 (DirectML) |
0.50 / 0.72 / 0.59 |
6/6 · 6/8 |
6/7 · 6/6 |
5/5 |
| INT8 |
0.49 / 0.73 / 0.59 |
7/7 · 7/8 |
6/8 · 6/6 |
5/5 |
- P is precision (correct / predicted) and R is recall (found / expected).
- Danbooru tag lists are often incomplete, so the general precision shown is a lower bound.
- Ground-truth tags that are not in the model's vocabulary are excluded (3 general, 1 character, 1 copyright).
- With 5 images, a difference of one tag between variants is within noise.
Preprocessing
- Resize so the longest side is 1008 px (bicubic), keeping aspect ratio.
- Paste centered on a 1008 × 1008 black canvas.
- Scale to [0, 1], then normalize with mean 0.5 and std 0.5.
Categories and thresholds
category |
Name |
Threshold |
| 0 |
general |
0.17 |
| 1 |
style |
0.15 |
| 3 |
copyright |
0.24 |
| 4 |
character |
0.27 |
| 5 |
meta |
0.17 |
| 9 |
rating |
0.41 |
Thresholds are the ones recommended in the original model card.
Usage
pip install onnxruntime pillow numpy huggingface_hub
# DirectML (Windows): pip install onnxruntime-directml pillow numpy huggingface_hub
import csv
import numpy as np
import onnxruntime as ort
from PIL import Image
from huggingface_hub import hf_hub_download
REPO = "Mexes/pixai-tagger-v1.0-onnx-fp32-fp16-int8"
CATEGORIES = {0: "general", 1: "style", 3: "copyright", 4: "character", 5: "meta", 9: "rating"}
THRESHOLDS = {"general": 0.17, "style": 0.15, "copyright": 0.24, "character": 0.27, "meta": 0.17, "rating": 0.41}
with open(hf_hub_download(REPO, "selected_tags.csv"), encoding="utf-8") as f:
tags = [(row["name"], CATEGORIES[int(row["category"])]) for row in csv.DictReader(f)]
def preprocess(path, size=1008):
img = Image.open(path).convert("RGB")
scale = size / max(img.size)
w, h = round(img.width * scale), round(img.height * scale)
img = img.resize((w, h), Image.Resampling.BICUBIC)
canvas = Image.new("RGB", (size, size))
canvas.paste(img, ((size - w) // 2, (size - h) // 2))
x = (np.asarray(canvas, dtype=np.float32) / 255.0 - 0.5) / 0.5
return x.transpose(2, 0, 1)[None]
session = ort.InferenceSession(
hf_hub_download(REPO, "model_fp16.onnx"),
providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
)
probs = session.run(["prediction"], {"input": preprocess("image.jpg")})[0][0]
results = {}
for (name, category), p in zip(tags, probs):
if p >= THRESHOLDS[category]:
results.setdefault(category, {})[name] = float(p)
print(results)
DirectML
so = ort.SessionOptions()
so.enable_mem_pattern = False
so.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
session = ort.InferenceSession(
"model_fp16.onnx", so,
providers=[("DmlExecutionProvider", {"device_id": 0}), "CPUExecutionProvider"],
)
Set device_id explicitly on systems with more than one GPU; without it the integrated GPU may be selected
(21.8 s vs 5.3 s per image in the test above).
Conversion
pip install torch safetensors numpy onnx onnxruntime onnxconverter-common
python scripts/export_pixai_v1_onnx.py out
python scripts/convert_pixai_v1_fp16.py --src out/model.onnx
python scripts/quantize_pixai_v1_int8.py --src out/model.onnx
The rotary position embeddings are rewritten with real-valued cos/sin, since the complex-number ops used by
the original implementation cannot be exported to ONNX.
License
Apache-2.0, same as the original model.
