Haruka-Tagger-beta-v6

What's New?

If you have been following my haruka-tagger adventures, you probably know that I released haruka-tagger-alpha.

I'm calling this beta. Because for 2 reasons:

1. I think I'm good / had enough of training haruka-tagger.
2. The model is quite good.

As a general recap, most danbooru taggers cap their vocab to around 10-13K tags. While it's fine for most cases, sometimes you really want a model to just say: "Why not just predict on everything lol"

This is haruka-tagger.

2-Level Coarse -> Fine Heads

Standard single-query attention pooling (num_queries=1) crushes all 32×32 = 1,024 spatial tokens into a single 768-d vector before classification. On multi-character scenes, cosplay mashups, or detailed backgrounds, that single vector suffers from feature dilution and cross-category interference.

In beta v6, we replaced flat pooling with a Coarse to Fine branching query tree. 6 Coarse Heads paired with 3 Child leaves which gives 18 combined queries.

Additionally, each of the leaves has it's own orthogonal projection matrix that preserves the full rank of the 768-d backbone.

Slot Sigmoid Routing

Since softmax(dim=-2)'d queries compete for patches, we used sigmoid(logits + bias) instead followed by patch-mass scaling M_q / (M_q + 4.0) across LayerNorm so empty queries can stay silent instead of amplifying noise.

Training Adjustments (Two-Sided IRFS + Per-Tag Loss Weighting)

We adjusted the training by lowered the irfs power. Basically instead of some tags being repeat for 4x, we scaled it down to 2x.

Per-Tag Loss Weighting was added (taking from animetimm lmao)

Evaluation

Evaluations are hard. Not only there is no standard across all models, additionally those that look "good" on paper tends to have oddities.

We used new-dan crawled by one of our friends. Since haruka-tagger is trained exclusively on images before/on 2026-01-01, all the images in new-dan are unseen images for haruka-tagger. That being said though... Other models may include images after 2026-01-01.

If you disagree with these results, we recommend to try out the model yourself.

Thresholds

Quickstart Usage

Option 1: Standalone PyTorch + safetensors

import json
import math
import numpy as np
import timm
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from safetensors.torch import load_file

IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32).reshape(3, 1, 1)
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32).reshape(3, 1, 1)


def preprocess_image(img: Image.Image, target_size: int = 448) -> torch.Tensor:
    """Alpha-composites onto white (255, 255, 255), letterbox-resizes (BICUBIC) to 448x448, and normalizes."""
    if img.mode != "RGB":
        img = img.convert("RGBA")
        bg = Image.new("RGBA", img.size, (255, 255, 255, 255))
        bg.alpha_composite(img)
        img = bg.convert("RGB")

    w, h = img.size
    if w != target_size or h != target_size:
        scale = min(target_size / w, target_size / h)
        nw, nh = max(1, int(w * scale)), max(1, int(h * scale))
        img = img.resize((nw, nh), Image.Resampling.BICUBIC)
        canvas = Image.new("RGB", (target_size, target_size), (255, 255, 255))
        canvas.paste(img, ((target_size - nw) // 2, ((target_size - nh) // 2)))
        img = canvas

    arr = np.asarray(img, dtype=np.float32).transpose(2, 0, 1) / 255.0
    arr = (arr - IMAGENET_MEAN) / IMAGENET_STD
    return torch.from_numpy(arr).unsqueeze(0)


class MHAttnPool(nn.Module):
    def __init__(self, dim: int = 768, num_heads: int = 16, num_queries: int = 18, leaf_dim: int = 64, pool_mode: str = "slot_sigmoid"):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.num_queries = num_queries
        self.pool_mode = pool_mode
        self.norm = nn.LayerNorm(dim)
        self.kv = nn.Linear(dim, dim * 2)
        self.num_coarse = min(6, num_queries)
        self.branches = num_queries // self.num_coarse
        self.out_dim = num_queries * leaf_dim
        self.q_coarse = nn.Parameter(torch.zeros(1, self.num_coarse, dim))
        self.trunk_norm = nn.LayerNorm(dim)
        self.q_child = nn.Parameter(torch.zeros(1, self.num_coarse, self.branches, dim))
        self.child_q_norm = nn.LayerNorm(dim)
        self.leaf_norm = nn.LayerNorm(dim)
        if pool_mode == "slot_sigmoid":
            self.q_coarse_bias = nn.Parameter(torch.full((1, 1, self.num_coarse, 1), -2.0))
            self.q_child_bias = nn.Parameter(torch.full((1, 1, num_queries, 1), -2.0))
        self.leaf_proj_weight = nn.Parameter(torch.empty(num_queries, dim, leaf_dim))
        self.leaf_proj_bias = nn.Parameter(torch.zeros(num_queries, leaf_dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, N, D = x.shape
        x = self.norm(x)
        k, v = self.kv(x).chunk(2, dim=-1)
        k = k.view(B, N, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, N, self.num_heads, self.head_dim).transpose(1, 2)
        scale = self.head_dim ** -0.5

        q_c = self.q_coarse.expand(B, self.num_coarse, D).view(B, self.num_coarse, self.num_heads, self.head_dim).transpose(1, 2)
        logits_c = (q_c @ k.transpose(-2, -1)) * scale
        attn_c = (
            torch.softmax(logits_c, dim=-2) * torch.sigmoid(logits_c + self.q_coarse_bias)
            if self.pool_mode == "slot_sigmoid"
            else torch.softmax(logits_c, dim=-1)
        )
        z_coarse = self.trunk_norm((attn_c @ v).transpose(1, 2).contiguous().view(B, self.num_coarse, D))

        q_leaf = self.child_q_norm((z_coarse.unsqueeze(2) + self.q_child).view(B, self.num_queries, D))
        q_l = q_leaf.view(B, self.num_queries, self.num_heads, self.head_dim).transpose(1, 2)
        logits_l = (q_l @ k.transpose(-2, -1)) * scale
        attn_l = (
            torch.softmax(logits_l, dim=-2) * torch.sigmoid(logits_l + self.q_child_bias)
            if self.pool_mode == "slot_sigmoid"
            else torch.softmax(logits_l, dim=-1)
        )
        z_child = (attn_l @ v).transpose(1, 2).contiguous().view(B, self.num_coarse, self.branches, D)

        z_tree = self.leaf_norm((z_coarse.unsqueeze(2) + z_child).view(B, self.num_queries, D))
        if self.pool_mode == "slot_sigmoid":
            mass_q = attn_l.sum(dim=-1).mean(dim=1).unsqueeze(-1)
            z_tree = z_tree * (mass_q / (mass_q + 4.0))

        leaf_vecs = torch.einsum("bqd,qdf->bqf", z_tree, self.leaf_proj_weight) + self.leaf_proj_bias
        return leaf_vecs.reshape(B, self.out_dim)


class HarukaStandalone(nn.Module):
    def __init__(self, meta: dict):
        super().__init__()
        self.backbone = timm.create_model(meta["backbone"], pretrained=False, num_classes=0, global_pool="")
        self.head_pool = MHAttnPool(
            dim=self.backbone.num_features,
            num_queries=meta.get("num_queries", 18),
            pool_mode=meta.get("pool_mode", "slot_sigmoid"),
        )
        self.head = nn.Linear(self.head_pool.out_dim, meta["total_tags"])

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        feats = self.backbone.forward_features(x)
        return self.head(self.head_pool(feats))


# Load metadata & weights
with open("haruka-tagger_metadata.json", "r", encoding="utf-8") as f:
    meta = json.load(f)

device = "cuda" if torch.cuda.is_available() else "cpu"
model = HarukaStandalone(meta)
model.load_state_dict(load_file("haruka-tagger.safetensors"), strict=True)
model.to(device).eval()

# Run inference on an image
x = preprocess_image(Image.open("example.jpg"), target_size=meta["img_size"]).to(device)
with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=(device == "cuda")):
    probs = torch.sigmoid(model(x).float())[0].cpu().numpy()

thresholds = {"character": 0.35, "copyright": 0.35, "artist": 0.25, "general": 0.35, "rating": 0.50, "year": 0.25}
idx_to_tag = meta["idx_to_tag"]
tag_to_cat = meta["tag_to_category"]

results = {c: [] for c in thresholds}
for idx in np.where(probs >= 0.15)[0]:
    tag = idx_to_tag[str(int(idx))]
    cat = tag_to_cat.get(tag, "general")
    score = float(probs[idx])
    if score >= thresholds.get(cat, 0.35):
        results[cat].append((tag, round(score, 4)))

for cat, items in results.items():
    items.sort(key=lambda x: -x[1])
    print(f"{cat.upper():<10}:", ", ".join(f"{t} ({s:.0%})" for t, s in items[:15]))

---

Option 2: ONNX Runtime (haruka-tagger.onnx)

import json
import numpy as np
import onnxruntime as ort
from PIL import Image

with open("haruka-tagger_metadata.json", "r", encoding="utf-8") as f:
    meta = json.load(f)

sess = ort.InferenceSession("haruka-tagger.onnx", providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
x_np = preprocess_image(Image.open("example.jpg"), target_size=meta["img_size"]).numpy()

logits = sess.run(["logits"], {"images": x_np})[0][0]
probs = 1.0 / (1.0 + np.exp(-logits))