PixAI Tagger v1.0 — ONNX (FP32 / FP16 / INT8)

Unofficial ONNX export of pixai-labs/pixai-tagger-v1.0 in FP32, FP16 and INT8. Runs with ONNX Runtime only.

Files

• 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)

Against Danbooru tags (5 posts)

• 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

1. Resize so the longest side is 1008 px (bicubic), keeping aspect ratio.
2. Paste centered on a 1008 × 1008 black canvas.
3. Scale to [0, 1], then normalize with mean 0.5 and std 0.5.

Categories and thresholds

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.