PixAI Tagger v1.0 — ONNX (FP16 + FP32) at 1008 / 672 / 336, zero-shot and fine-tuned (+ 448, Hammy-only)

This version is primarily assembled for Hammy, a Chinese-language video organizing & archiving tool. The 448-ft model (448-ft/) is included for Hammy only: it is calibrated for character/IP tagging (0.87 / 0.895 in its config.json); the other categories were not tuned for it.

Community conversion of pixai-labs/pixai-tagger-v1.0 (Apache-2.0) to ONNX in FP16 and FP32, plus a script that builds TensorRT engines from them on any GPU.

For the full model documentation (architecture, training data, original usage), see the upstream model card.

Two families ship side by side:

* general / character / copyright. The fine-tuned graphs must be used with their own re-swept thresholds — they are stricter (see Accuracy); the root config.json belongs to the root graphs only. Each fine-tuned onnx/ folder is self-contained (graph + .data + config.json + preprocessor_config.json + tag_zh.json), so it can be copied out on its own; the runners pick the config up automatically (Tagger(..., ft=True)). If you re-tune thresholds by hand, update both the onnx/config.json (what the runners read) and the safetensors/config.json (what from_pretrained reads).

Speed

Forward time (median, per image; batch 1 unless noted) — RTX 5070 (12 GB), Windows, TensorRT 11.3 / onnxruntime-gpu 1.30:

• On Tensor-Core cards FP16 is the default; FP32 compute is ~4× slower there.
• GTX 16xx inverts the ranking. The 1660 Super is sm_75 without Tensor Cores, so fp16 GEMM kernels fall back to slow SIMT paths (4698 ms @1008 vs the FP32 engine's 1448 ms and the PyTorch checkpoint's 1789 ms) — there only the FP32 engine beats the PyTorch checkpoint. Build with --fp32 on those cards; the query-chunked FP32 engine also fits their 6 GB of VRAM.
• 336-ft numbers are per image at batch 2 (rows 1008/672 are batch 1).

Files

.gitattributes routes *.onnx, *.onnx.data, *.safetensors through LFS; the FP32 graphs need their .data sidecar next to them (relative filename, resolved by ORT and by TensorRT).

Accuracy

Method. The oracle is the original FP32 checkpoint running at 1008 (PyTorch, no autocast, batch 1). Tags are compared tag-by-tag above each model's own recommended per-category thresholds (table below) on 1,000 real test images (WebP, long edge 500–2374) prepared independently of the fine-tuning data. A "flip" counts a tag crossing its threshold in either direction; focus = character + copyright — the identity-critical tags (~1,700 of them above threshold at 1008); the 15k+ general tags are dense around their threshold by design and are the low-stakes category. zs = the original weights with the pos/RoPE grids rebuilt for the target grid (no training); ft = distilled from the oracle on 103,821 images (2 epochs).

The fine-tuned thresholds were re-swept on 3,000 held-out images; they are stricter and must be used with the fine-tuned graphs (the runners pick them up from the config next to the graph).

1008 — the conversion is (almost) lossless

Focus-tag flips vs the oracle, per runtime:

max |Δp| vs the checkpoint: 7.1e-3 (FP32) / 1.1e-2 (FP16), mean |Δp| ~1e-6 on both; every flipped tag sits right at its threshold (median baseline 0.170 vs the 0.17 general threshold, median movement 0.002, max ≤0.01). Safe default.

672 / 336 — what fine-tuning buys

Focus-tag flips vs the oracle, fp32 (TensorRT engine; the onnxruntime and fp16 runs land within a few tags of these numbers, see Runtime consistency):

• 672-ft: focus losses −13% (233→203), focus false positives −59% (238→97), general flips −38% (15,652→9,667).
• 336-ft: losses −19% (432→350), false positives −73% (530→143), general flips −47% (24,896→13,170).
• The zs rows are what converting to a smaller grid costs without re-training (e.g. yae_miko 0.997 @1008 → 0.0002 @672; threshold_shift does not recover such losses — they land far below any shift worth applying) — they are the reason the fine-tuned checkpoints exist.

Runtime consistency

The same model across the four runtimes (TRT/ORT × fp16/fp32) agrees closely: focus flips stay within ±3 tags and the dense general category within ±13. fp16 vs fp32 is identical on identity tags, so the differences above come from resolution and fine-tuning, not from the runtime or precision. fp16 is recommended on GPU; on GTX 16xx use fp32 (see Speed & memory).

INT8 (CPU)

192-image subset, int8_dyn.onnx vs the fp32 ORT path:

CPU-only (on GPU the quantized ops fall back to CPU kernels, ~3 img/s); ~1.6× the FP32 speed on CPU at ¼ the size. Must be used with the ft336 thresholds.

Notes

• ORT CUDA FP16 DLL trap: in a process that has imported both torch and onnxruntime, manually preloading the pip nvidia-cudnn-cu12 / nvidia-cublas-cu12 DLLs on top of torch's bundled cuDNN makes ORT silently return wrong results (reproduced during this testing). Let ORT resolve naturally (torch first, as both runners do), or run ORT in a torch-free process.
• The test set is independent of the training data; the fine-tuned models share the training pool's domain (anime illustrations) — out-of-domain content (photos, manga pages) was not evaluated.
• Source-resolution robustness (measured on the zero-shot conversion): pre-resizing large sources to long edge 1008 costs only borderline tag flips and zero confident losses; even 672-px sources upscaled into the 1008 pipeline lose fewer tags than running the 672 grid — identity recognition is far more sensitive to the model's grid size than to the source image size.

Rule of thumb: identity-critical tagging → 1008; bulk/backlog tagging → 336-ft (batch 2, fastest) or 672-ft when identity matters more and throughput less.

Usage

# Option A — onnxruntime (any CUDA GPU)
from tagger_ort import Tagger
tagger = Tagger(size=672)                       # zero-shot root graphs, FP16
tagger = Tagger(size=672, precision="fp32")     # zero-shot FP32 graph
tagger = Tagger(size=336, ft=True)              # fine-tuned: 336-ft/onnx/fp16.onnx (batch 2)
tagger = Tagger(size=448, ft=True)              # 448-ft/onnx/fp16.onnx (batch 1, Hammy-only model)
tagger = Tagger(size=672, ft=True, precision="fp32")
r = tagger.tag("image.png")                     # {"character": {...}, "copyright": {...}}
r = tagger.tag(["a.png", "b.jpg", "c.png"])     # batched; 336 is fed two images at a time

# Option B — TensorRT (fastest): build an engine for your GPU first (under a minute):
#   python export_trt.py build 672                           (FP16, from tagger_672_fp16.onnx)
#   python export_trt.py build --onnx 336-ft/onnx/fp16.onnx  (-> 336-ft/onnx/fp16.trt, batch 2)
#   --workspace caps device scratch in GiB (default 4) — use 2 on <=6 GB cards.
from tagger_trt import Tagger
tagger = Tagger(size=336, ft=True)              # loads 336-ft/onnx/fp16.trt
tagger = Tagger(size=672, ft=True, precision="fp32")

• Pass model_dir=<folder> to point at any folder of your own: a folder holding tagger_<size>_<tag>.onnx (+ config.json) works, and so does a -ft-style folder (onnx/<tag>.onnx + onnx/config.json, checkpoint alongside in safetensors/). ft=True is just shorthand for the repo's 336-ft/ / 672-ft/ folders.
• Thresholds per category come from the config.json next to the graph (root graphs: repo root; fine-tuned: <size>-ft/onnx/config.json); pass threshold_shift=0.05 to prune borderline tags.
• Build FP32 engines only from the shipped fp32.onnx graphs (dynamo-exported, with query-chunked global attention so the engine fits 6 GB cards): TensorRT silently miscompiles a legacy-TorchScript FP32 graph of this model, so never substitute one.
• Engines are not portable across GPU architectures or TensorRT versions — always rebuild on the target machine, and expect two builds from the same ONNX to differ by ~1e-3 in probability (never compare engines bit-for-bit).

Optional: Chinese tag names

tag_zh.json ships next to config.json (in the fine-tuned folders it sits in onnx/ with the graphs; the repo root copy is the fallback) and contains Chinese names for the character and copyright tags only (no general tags), built for Hammy; translation quality is not guaranteed. Both runners ship with translation off by default — pass zh=True to enable it (zh=True requires the file to exist):

tagger = Tagger(size=672, zh=True)

{"character": {"初音未来": 1.0}, "copyright": {"VOCALOID": 1.0}}

Requirements

Pick one path, not all of them:

# A — onnxruntime only:
#     pip install onnxruntime-gpu pillow numpy nvidia-cudnn-cu12 nvidia-cublas-cu12

# B — TensorRT engine, built on this machine from the shipped ONNX:
#     pip install tensorrt nvidia-cuda-runtime nvidia-cuda-nvrtc --extra-index-url https://pypi.nvidia.cn
#     (`tensorrt` itself has no win_amd64 wheel on PyPI — only NVIDIA's index carries the
#      ~2 GB of binaries. `tensorrt` 11.3 unconditionally pulls `tensorrt_cu13`, a CUDA 13
#      runtime — the driver must support CUDA 13 (>= 580). The wheels do NOT bundle
#      cudart/nvrtc: nvinfer_11.dll imports nvcudart_hybrid64.dll and nvrtc64_130_0.dll,
#      hence the two extra packages. Tested with TensorRT 11.3 on Windows; TRT 11
#      supports sm_75 and up.)

# C — B, plus running the engine (tagger_trt.py imports the repo's tagger_pipeline.py):
#     pip install torch --index-url https://download.pytorch.org/whl/cu128
#     pip install transformers timm torchvision pillow

# Only `export_trt.py export <size>` needs a checkpoint (the original repo, or a
# `*-ft/safetensors` folder via --model-dir); `build` and both runners run from this
# repo alone.

INT8 needs onnxruntime CPU: pip install onnxruntime pillow numpy.

License & attribution

Converted from pixai-labs/pixai-tagger-v1.0 (© PixAI labs, Apache-2.0). The fine-tuned 336/448/672 checkpoints in *-ft/ were distilled from that model's own outputs; this repository redistributes converted weights and runner scripts under Apache-2.0; all model credit belongs to the original authors.