"""DINOv3 ViT-H/16+ Tagger — Fully Standalone Inference Script Zero dependency on transformers, trainer code, or any internal module. Only requires: torch, torchvision, safetensors, Pillow, requests. pip install torch torchvision safetensors Pillow requests The DINOv3 ViT-H/16+ architecture is implemented directly here, with weights loaded from a .safetensors checkpoint. The state-dict key names match the HuggingFace transformers layout exactly so checkpoints are interchangeable. Usage ----- # Single image, top-30 tags: python inference_tagger_standalone.py \ --checkpoint tagger_checkpoints/2026-03-28_22-57-47.safetensors \ --vocab tagger_vocab.json \ --images photo.jpg \ --topk 30 # URL input: python inference_tagger_standalone.py \ --checkpoint tagger_checkpoints/2026-03-28_22-57-47.safetensors \ --vocab tagger_vocab.json \ --images https://example.com/photo.jpg # Threshold instead of top-k: python inference_tagger_standalone.py ... --threshold 0.4 # Pipe-friendly comma-separated tags (one line per image): python inference_tagger_standalone.py ... --format tags # JSON output: python inference_tagger_standalone.py ... --format json Output formats (--format) ------------------------- pretty (default) — human-readable table with scores tags — comma-separated tag string, one line per image json — JSON array of {file, tags: [{tag, score}]} objects """ from __future__ import annotations import argparse import json import math import sys from functools import lru_cache from io import BytesIO from pathlib import Path import requests import torch import torch.nn as nn import torch.nn.functional as F import torchvision.transforms.v2 as v2 from PIL import Image from safetensors.torch import load_file # ============================================================================= # DINOv3 ViT-H/16+ — hardcoded architecture # All hyperparameters match facebook/dinov3-vith16plus-pretrain-lvd1689m # ============================================================================= D_MODEL = 1280 N_HEADS = 20 HEAD_DIM = D_MODEL // N_HEADS # 64 N_LAYERS = 32 D_FFN = 5120 N_REGISTERS = 4 PATCH_SIZE = 16 ROPE_THETA = 100.0 ROPE_RESCALE = 2.0 # pos_embed_rescale applied at inference LN_EPS = 1e-5 LAYERSCALE = 1.0 # --------------------------------------------------------------------------- # RoPE helpers # --------------------------------------------------------------------------- @lru_cache(maxsize=32) def _patch_coords_cached(h: int, w: int, device_str: str) -> torch.Tensor: """Normalised [-1,+1] patch-centre coordinates (float32, cached).""" device = torch.device(device_str) cy = torch.arange(0.5, h, dtype=torch.float32, device=device) / h cx = torch.arange(0.5, w, dtype=torch.float32, device=device) / w coords = torch.stack(torch.meshgrid(cy, cx, indexing="ij"), dim=-1).flatten(0, 1) coords = 2.0 * coords - 1.0 # [0,1] → [-1,+1] coords = coords * ROPE_RESCALE return coords # [h*w, 2] def _build_rope(h_patches: int, w_patches: int, dtype: torch.dtype, device: torch.device): """Return (cos, sin) of shape [1, 1, h*w, HEAD_DIM] for broadcasting.""" coords = _patch_coords_cached(h_patches, w_patches, str(device)) # [P, 2] inv_freq = 1.0 / (ROPE_THETA ** torch.arange( 0, 1, 4 / HEAD_DIM, dtype=torch.float32, device=device)) # [D/4] angles = 2 * math.pi * coords[:, :, None] * inv_freq[None, None, :] # [P, 2, D/4] angles = angles.flatten(1, 2).tile(2) # [P, D] cos = torch.cos(angles).to(dtype).unsqueeze(0).unsqueeze(0) # [1,1,P,D] sin = torch.sin(angles).to(dtype).unsqueeze(0).unsqueeze(0) return cos, sin def _rotate_half(x: torch.Tensor) -> torch.Tensor: h = x.shape[-1] // 2 return torch.cat((-x[..., h:], x[..., :h]), dim=-1) def _apply_rope(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor): """Apply RoPE only to patch tokens (skip CLS + register prefix).""" n_pre = 1 + N_REGISTERS q_pre, q_pat = q[..., :n_pre, :], q[..., n_pre:, :] k_pre, k_pat = k[..., :n_pre, :], k[..., n_pre:, :] q_pat = q_pat * cos + _rotate_half(q_pat) * sin k_pat = k_pat * cos + _rotate_half(k_pat) * sin return torch.cat([q_pre, q_pat], dim=-2), torch.cat([k_pre, k_pat], dim=-2) # --------------------------------------------------------------------------- # Building blocks # --------------------------------------------------------------------------- class _Attention(nn.Module): def __init__(self): super().__init__() self.q_proj = nn.Linear(D_MODEL, D_MODEL, bias=True) self.k_proj = nn.Linear(D_MODEL, D_MODEL, bias=False) self.v_proj = nn.Linear(D_MODEL, D_MODEL, bias=True) self.o_proj = nn.Linear(D_MODEL, D_MODEL, bias=True) def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: B, S, _ = x.shape q = self.q_proj(x).view(B, S, N_HEADS, HEAD_DIM).transpose(1, 2) k = self.k_proj(x).view(B, S, N_HEADS, HEAD_DIM).transpose(1, 2) v = self.v_proj(x).view(B, S, N_HEADS, HEAD_DIM).transpose(1, 2) q, k = _apply_rope(q, k, cos, sin) out = F.scaled_dot_product_attention(q, k, v, scale=HEAD_DIM ** -0.5) return self.o_proj(out.transpose(1, 2).reshape(B, S, D_MODEL)) class _GatedMLP(nn.Module): def __init__(self): super().__init__() self.gate_proj = nn.Linear(D_MODEL, D_FFN, bias=True) self.up_proj = nn.Linear(D_MODEL, D_FFN, bias=True) self.down_proj = nn.Linear(D_FFN, D_MODEL, bias=True) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) class _Block(nn.Module): def __init__(self): super().__init__() self.norm1 = nn.LayerNorm(D_MODEL, eps=LN_EPS) self.attention = _Attention() self.layer_scale1 = nn.Parameter(torch.full((D_MODEL,), LAYERSCALE)) self.norm2 = nn.LayerNorm(D_MODEL, eps=LN_EPS) self.mlp = _GatedMLP() self.layer_scale2 = nn.Parameter(torch.full((D_MODEL,), LAYERSCALE)) def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: x = x + self.attention(self.norm1(x), cos, sin) * self.layer_scale1 x = x + self.mlp(self.norm2(x)) * self.layer_scale2 return x # --------------------------------------------------------------------------- # Full backbone # --------------------------------------------------------------------------- class DINOv3ViTH(nn.Module): """DINOv3 ViT-H/16+ backbone. Accepts any H, W that are multiples of 16. Returns last_hidden_state [B, 1+R+P, D_MODEL]. Token layout: [CLS, reg_0..reg_3, patch_0..patch_N]. State-dict keys are intentionally identical to the HuggingFace transformers layout so .safetensors checkpoints load without remapping. """ def __init__(self): super().__init__() # These names must match HF exactly self.embeddings = _Embeddings() self.layer = nn.ModuleList([_Block() for _ in range(N_LAYERS)]) self.norm = nn.LayerNorm(D_MODEL, eps=LN_EPS) def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): # HF stores layer_scale as a sub-module with a "lambda1" parameter; # we store it as a plain Parameter directly on _Block. # Remap "layer.i.layer_scale{1,2}.lambda1" → "layer.i.layer_scale{1,2}" for k in list(state_dict.keys()): if k.startswith(prefix) and ".layer_scale" in k and k.endswith(".lambda1"): new_k = k[:-len(".lambda1")] state_dict[new_k] = state_dict.pop(k) # Drop rope_embeddings buffer (computed on-the-fly) for k in list(state_dict.keys()): if k.startswith(prefix) and "rope_embeddings" in k: state_dict.pop(k) super()._load_from_state_dict( state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: B, _, H, W = pixel_values.shape x = self.embeddings(pixel_values) # [B, 1+R+P, D] h_p, w_p = H // PATCH_SIZE, W // PATCH_SIZE cos, sin = _build_rope(h_p, w_p, x.dtype, pixel_values.device) for block in self.layer: x = block(x, cos, sin) return self.norm(x) class _Embeddings(nn.Module): """Patch + CLS + register token embeddings. Key names match HF: embeddings.cls_token, embeddings.register_tokens, embeddings.patch_embeddings.{weight,bias}. """ def __init__(self): super().__init__() self.cls_token = nn.Parameter(torch.empty(1, 1, D_MODEL)) self.mask_token = nn.Parameter(torch.zeros(1, 1, D_MODEL)) # unused at inference self.register_tokens = nn.Parameter(torch.empty(1, N_REGISTERS, D_MODEL)) self.patch_embeddings = nn.Conv2d(3, D_MODEL, kernel_size=PATCH_SIZE, stride=PATCH_SIZE) def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: B = pixel_values.shape[0] dtype = self.patch_embeddings.weight.dtype patches = self.patch_embeddings(pixel_values.to(dtype)).flatten(2).transpose(1, 2) cls = self.cls_token.expand(B, -1, -1) regs = self.register_tokens.expand(B, -1, -1) return torch.cat([cls, regs, patches], dim=1) # ============================================================================= # Tagger head # ============================================================================= class DINOv3Tagger(nn.Module): """DINOv3 ViT-H/16+ backbone + linear projection head. features = concat(CLS, reg_0..reg_3) → [B, (1+R)*D] projection: Linear → [B, num_tags] """ def __init__(self, num_tags: int, projection_bias: bool = False): super().__init__() self.backbone = DINOv3ViTH() self.projection = nn.Linear((1 + N_REGISTERS) * D_MODEL, num_tags, bias=projection_bias) def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: hidden = self.backbone(pixel_values) # [B, S, D] cls = hidden[:, 0, :] # [B, D] regs = hidden[:, 1: 1 + N_REGISTERS, :].flatten(1) # [B, R*D] features = torch.cat([cls, regs], dim=-1) # [B, (1+R)*D] return self.projection(features.float()) # fp32 for stability # ============================================================================= # Image preprocessing # ============================================================================= _IMAGENET_MEAN = [0.485, 0.456, 0.406] _IMAGENET_STD = [0.229, 0.224, 0.225] def _snap(x: int, m: int) -> int: return max(m, (x // m) * m) def _open_image(source) -> Image.Image: s = str(source) if s.startswith("http://") or s.startswith("https://"): r = requests.get(s, timeout=30) r.raise_for_status() return Image.open(BytesIO(r.content)).convert("RGB") return Image.open(source).convert("RGB") def preprocess_image(source, max_size: int = 1024) -> torch.Tensor: """Load and preprocess an image → [1, 3, H, W] float32, ImageNet-normalised.""" img = _open_image(source) w, h = img.size scale = min(1.0, max_size / max(w, h)) new_w = _snap(round(w * scale), PATCH_SIZE) new_h = _snap(round(h * scale), PATCH_SIZE) return v2.Compose([ v2.Resize((new_h, new_w), interpolation=v2.InterpolationMode.LANCZOS), v2.ToImage(), v2.ToDtype(torch.float32, scale=True), v2.Normalize(mean=_IMAGENET_MEAN, std=_IMAGENET_STD), ])(img).unsqueeze(0) # ============================================================================= # Tagger wrapper # ============================================================================= class Tagger: """Inference wrapper for DINOv3Tagger (ViT-H/16+). Parameters ---------- checkpoint_path : str Path to a .safetensors or .pth checkpoint saved by TaggerTrainer. vocab_path : str Path to tagger_vocab.json ({"idx2tag": [...]}). device : str "cuda", "cuda:0", "cpu", etc. dtype : torch.dtype bfloat16 recommended on Ampere+; float16 for older GPUs; float32 for CPU. max_size : int Long-edge cap in pixels before feeding to the model. """ def __init__( self, checkpoint_path: str, vocab_path: str, device: str = "cuda", dtype: torch.dtype = torch.bfloat16, max_size: int = 1024, ): self.device = torch.device(device if torch.cuda.is_available() or device == "cpu" else "cpu") self.dtype = dtype self.max_size = max_size with open(vocab_path) as f: data = json.load(f) self.idx2tag: list[str] = data["idx2tag"] self.num_tags = len(self.idx2tag) print(f"[Tagger] Vocabulary: {self.num_tags:,} tags") self.model = DINOv3Tagger(num_tags=self.num_tags) print(f"[Tagger] Loading checkpoint: {checkpoint_path}") if checkpoint_path.endswith((".safetensors", ".sft")): sd = load_file(checkpoint_path, device=str(self.device)) else: sd = torch.load(checkpoint_path, map_location=str(self.device)) missing, unexpected = self.model.load_state_dict(sd, strict=False, assign=True) if missing: print(f"[Tagger] Missing keys ({len(missing)}): {missing[:5]}{'...' if len(missing) > 5 else ''}") if unexpected: print(f"[Tagger] Unexpected keys ({len(unexpected)}): {unexpected[:5]}{'...' if len(unexpected) > 5 else ''}") self.model.backbone = self.model.backbone.to(dtype=dtype) self.model = self.model.to(self.device) self.model.eval() print(f"[Tagger] Ready on {self.device} ({dtype})") @torch.no_grad() def predict(self, image, topk: int | None = 30, threshold: float | None = None) -> list[tuple[str, float]]: """Tag a single image (local path or URL). Specify either topk OR threshold. Returns [(tag, score), ...] desc.""" if topk is None and threshold is None: topk = 30 pv = preprocess_image(image, max_size=self.max_size).to(self.device) with torch.autocast(device_type=self.device.type, dtype=self.dtype): logits = self.model(pv)[0] scores = torch.sigmoid(logits.float()) if topk is not None: values, indices = scores.topk(min(topk, self.num_tags)) else: assert threshold is not None indices = (scores >= threshold).nonzero(as_tuple=True)[0] values = scores[indices] order = values.argsort(descending=True) indices, values = indices[order], values[order] return [(self.idx2tag[i], float(v)) for i, v in zip(indices.tolist(), values.tolist())] @torch.no_grad() def predict_batch(self, images, topk: int | None = 30, threshold: float | None = None) -> list[list[tuple[str, float]]]: """Tag multiple images (processed individually for mixed resolutions).""" return [self.predict(img, topk=topk, threshold=threshold) for img in images] # ============================================================================= # Output formatters # ============================================================================= def _fmt_pretty(path: str, results) -> str: lines = [f"\n{'─' * 60}", f" {path}", f"{'─' * 60}"] for rank, (tag, score) in enumerate(results, 1): bar = "█" * int(score * 20) lines.append(f" {rank:>3}. {score:.3f} {bar:<20} {tag}") return "\n".join(lines) def _fmt_tags(results) -> str: return ", ".join(tag for tag, _ in results) def _fmt_json(path: str, results) -> dict: return {"file": path, "tags": [{"tag": t, "score": round(s, 4)} for t, s in results]} # ============================================================================= # CLI # ============================================================================= def main(): parser = argparse.ArgumentParser( description="DINOv3 ViT-H/16+ tagger inference (standalone, no transformers dep)", formatter_class=argparse.RawDescriptionHelpFormatter, ) parser.add_argument("--checkpoint", required=True, help="Path to .safetensors or .pth checkpoint") parser.add_argument("--vocab", required=True, help="Path to tagger_vocab.json") parser.add_argument("--images", nargs="+", required=True, help="Image paths and/or http(s) URLs") parser.add_argument("--device", default="cuda", help="Device: cuda, cuda:0, cpu, … (default: cuda)") parser.add_argument("--max-size", type=int, default=1024, help="Long-edge cap in pixels, multiple of 16 (default: 1024)") mode = parser.add_mutually_exclusive_group() mode.add_argument("--topk", type=int, default=30, help="Return top-k tags (default: 30)") mode.add_argument("--threshold", type=float, help="Return all tags with score >= threshold") parser.add_argument("--format", choices=["pretty", "tags", "json"], default="pretty", help="Output format (default: pretty)") args = parser.parse_args() tagger = Tagger(checkpoint_path=args.checkpoint, vocab_path=args.vocab, device=args.device, max_size=args.max_size) topk, threshold = (None, args.threshold) if args.threshold else (args.topk, None) json_out = [] for src in args.images: is_url = str(src).startswith("http://") or str(src).startswith("https://") if not is_url and not Path(src).exists(): print(f"[warning] File not found: {src}", file=sys.stderr) continue results = tagger.predict(src, topk=topk, threshold=threshold) if args.format == "pretty": print(_fmt_pretty(src, results)) elif args.format == "tags": print(_fmt_tags(results)) elif args.format == "json": json_out.append(_fmt_json(src, results)) if args.format == "json": print(json.dumps(json_out, indent=2, ensure_ascii=False)) if __name__ == "__main__": main()