Haruka-Tagger (EVA02-Base-448, 98.8k Tags, 18-Query Slot-Sigmoid Tree)

An ultra-wide-vocabulary (98,782 tags) Danbooru image tagger co-trained end-to-end on an 86M-parameter eva02_base_patch14_448 backbone and a 2-level 18-query spatial attention tree.

Compute (1 Node of 8x H100) graciously sponsored by alpindale from dphn.ai.

---

Overview

Most community Danbooru taggers (WD-v3, AnimeTIMM, PixAI v0.9) cap their vocabulary around 10,000–13,500 tags—which works well for common clothing and popular characters, but drops artist attribution, niche copyrights, compositional micro-tags, and long-tail characters (45–500 posts).

Haruka-Tagger trains a compact eva02_base_patch14_448 (86.3M backbone) alongside a 2-Level 18-Query Slot-Sigmoid Attention Tree (MHAttnPool) from Step 0 across 9,253,992 wHash-deduplicated Danbooru images to predict 98,782 tags in a single forward pass at ~605 images/sec (batch_size=16, ~1.5 GB VRAM on H100):

Files in This Repository

• haruka-tagger.safetensors (809 MB) — Full PyTorch weights (backbone.*, head_pool.*, head.*).
• haruka-tagger.onnx (822 MB) — Standalone exported ONNX graph (input: images (B, 3, 448, 448) -> output: logits (B, 98782)).
• haruka-tagger_metadata.json (6.5 MB) — Complete tag index mapping (idx_to_tag, tag_to_category) and model configuration.
• haruka_infer.py — Self-contained CLI inference script, 2-level query tree inspector, and per-tag spatial attribution heatmap visualizer (python haruka_infer.py image.jpg --save-vis tree_vis.jpg).

---

Architecture & Training Recipe

1. Co-Trained 2-Level 18-Query Tree (MHAttnPool)

Standard single-query attention pooling (num_queries=1) forces the ViT backbone to squeeze all 32×32 = 1,024 spatial patch tokens into a single 768-d vector before classification. That forces the backbone's self-attention blocks to smear background, props, outfits, facial traits, and artist brushwork together globally across patches, causing feature dilution on multi-character scenes and cosplay/parody mashups.

Haruka-Tagger co-trains the backbone from Step 0 with a 2-Level Branching Query Tree (6 Coarse Trunks × 3 Child Leaves = 18 Queries):

• Pass 1 (6 Coarse Trunks): Six orthogonally initialized trunk queries scan the 32×32 patch grid to partition the canvas into coarse spatial/semantic regions.
• Pass 2 (18 Child Leaves): Three child queries branch off each parent trunk (z_coarse + q_child), attending to fine local details and inheriting their parent trunk via a residual connection (z_tree = LayerNorm(z_coarse + z_child)).
• Full-Rank Independent Per-Leaf Projections ((18, 768, 64)): Instead of a shared 768 -> 64 linear bottleneck (which collapses all 18 queries into the same 64-dimensional subspace), each leaf has its own independent (768, 64) projection matrix initialized orthogonally across $\mathbb{R}^{768}$, yielding a 18 × 64 = 1,152-d full-rank representation.
• Slot-Sigmoid Routing (slot_sigmoid) + Patch-Mass Scaling: Queries compete for patches via softmax(logits, dim=-2) while patches gate independently via sigmoid(logits + bias), followed by patch-mass scaling M_q / (M_q + 4.0) after LayerNorm so queries whose visual concept is absent can stay quiet instead of amplifying background noise.

2. Damped Two-Sided IRFS + Per-Tag Loss Weighting

Training across 98,782 tags spans five orders of magnitude in class support (1girl on 6.4M+ images down to tail artists/characters on 45 images). We combine both image-level and tag-level balancing:

• Damped Instance Repeat Factor Sampling (irfs_power = 0.5, [0.45x .. 2.0x]): Keeps 45-image long-tail tags alive in training batches without starving the backbone of everyday general-tag images.
• Per-Tag Loss Weighting ([0.40 .. 1.00]) inside AsymmetricLoss(gamma_pos=0.0, gamma_neg=2.0, clip=0.0): Scales down ultra-common tags (1girl = 0.40, solo = 0.41) relative to mid/rare tags (0.65–1.00) on the same image so common tags cannot hijack 80% of the backbone gradient when an image is sampled.
• Decoupled Sparse Classifier Weight Decay (1e-3 backbone / 1e-4 head) + Stochastic Depth (drop_path_rate = 0.2): Prevents AdamW decoupled weight decay from eroding rare classifier rows on the 99.9% of steps where a tail tag is absent from the batch.

3. Warmup–Stable–Decay (WSD) Training & Detail-Preserving Transforms

• No Destructive Cutout or Rotation Blur: Trained with clean 448×448 bicubic letterbox padding + horizontal flip (cutout = 0, rotation = 0, color_jitter = 0), preserving 1–3px micro-lineart details (@_@, eyelashes, piercing, white_outline, sweatdrop).
• WSD Schedule: Trained for 6 epochs at a stable high learning rate (lr = 3e-4, backbone_lr = 7.5e-5) to let the 18 queries and backbone organize globally, followed by a dedicated 3-epoch cosine cooldown (1.5e-4 -> 0) that jumped validation F1 to 68.23% (val_f1@35 = 63.16%).

---

Standalone Evaluation on Unseen 2026 Danbooru Uploads

Evaluated on 1,000 unseen post-cutoff 2026 Danbooru uploads (new-dan) across all 98,782 native tags (~605.4 img/sec at batch_size=16, 1,545 MB peak VRAM on H100). Rather than arguing over synthetic competitor scoreboards (since taggers target very different vocabulary sizes and resolutions), we report Haruka-Tagger's standalone metrics below and encourage you to test it directly on your own images:

General Tag Performance by Training Support Bucket

---

Sample 2-Level Query Tree Decompositions

Using haruka_infer.py, you can inspect which of the 18 leaves (6 Coarse Trunks × 3 Child Leaves) fired each tag:

Example 1: Character + Artist + Compositional Micro-Tags (Brave Witches by tama_kitsune)

• CHARACTER  ( 1): nikka_edvardine_katajainen (99.9%)
• COPYRIGHT  ( 2): brave_witches (91.5%), world_witches_series (85.0%)
• ARTIST     ( 1): tama_kitsune (98.7%)
• GENERAL    (47): 1girl (99.7%), breasts (99.7%), pantyhose (98.3%), animal_ears (98.0%), playboy_bunny (97.0%),
                   solo (95.4%), blush (93.8%), short_hair (92.8%), rabbit_ears (92.7%), tail (87.5%),
                   leotard (84.7%), large_breasts (84.5%), open_mouth (80.9%), rabbit_tail (80.8%), blonde_hair (78.7%),
                   looking_at_viewer (72%), detached_collar (70%), sweat (70%), blue_background (66%),
                   fake_animal_ears (65%), blue_leotard (65%), white_pantyhose (64%), embarrassed (58%),
                   covering_breasts (56%), border (54%), blue_bowtie (54%), white_outline (53%),
                   white_border (50%), @_@ (49%), fake_tail (48%), outside_border (41%), wrist_cuffs (37%)
  ├─ Leaf 0.0 (Q00) ──► playboy_bunny (97%), world_witches_series (85%), leotard (85%), covering_breasts (56%), white_outline (53%)
  ├─ Leaf 0.1 (Q01) ──► blush (94%), open_mouth (81%), sweat (70%), embarrassed (58%), sweatdrop (45%)
  ├─ Leaf 1.0 (Q03) ──► detached_collar (70%), blue_bowtie (54%), @_@ (49%)
  ├─ Leaf 1.2 (Q05) ──► nikka_edvardine_katajainen (100%), short_hair (93%), brave_witches (91%), white_pantyhose (64%)
  ├─ Leaf 3.2 (Q11) ──► bowtie (65%), border (54%), white_border (50%), outside_border (41%)
  └─ Leaf 4.1 (Q13) ──► tama_kitsune (99%)

Example 2: Parody / Crossover Meme (Minakami Mai from Nichijou Redrawn in the "Osaka" Azumanga Daioh Pose by denkirk)

• CHARACTER  ( 2): minakami_mai (99.7%), kasuga_ayumu (56.3%)
• COPYRIGHT  ( 2): nichijou (97.2%), azumanga_daioh (84.9%)
• ARTIST     ( 1): denkirk (40.3%)
• GENERAL    (14): 1girl (99.7%), solo (99.5%), glasses (97.6%), portrait (68.4%), long_hair (61.6%),
                   brown_hair (61.3%), expressionless (48.4%), closed_mouth (47.2%), black_hair (45.1%),
                   looking_at_viewer (43.3%), black_eyes (41.7%), brown_eyes (37.3%),
                   tokisadame_school_uniform (36.7%), simple_background (36.4%)

---

Quickstart Usage

Option 1: Included CLI + Tree Inspector & Spatial Heatmap Visualizer (haruka_infer.py)

python haruka_infer.py path/to/image.jpg --save-vis tree_vis.jpg

---

Option 2: Standalone PyTorch + safetensors

import json
import numpy as np
import timm
import torch
import torch.nn as nn
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 3: 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))