khasinski/siglip2-moderation-heads

SigLIP2 Moderation Heads

Lightweight MLP classification heads for image moderation, designed to run on top of frozen SigLIP2 embeddings.

Model Description

Three independent binary classifiers for detecting:

Each head is a tiny MLP (~197K params, <1MB) that takes a 768-dim SigLIP2 embedding and outputs a probability score.

Architecture

SigLIP2 Encoder (frozen, 400M params)
         │
         ▼ 768-dim embedding
         │
    ┌────┴────┐
    │ MLP Head │  ← Linear(768,256) → ReLU → Dropout(0.2) → Linear(256,1)
    └────┬────┘
         │
         ▼ logits → sigmoid → probability

Quick Start

Python (PyTorch)

import torch
from transformers import AutoModel, AutoProcessor
from PIL import Image

# Load SigLIP2 encoder
processor = AutoProcessor.from_pretrained("google/siglip2-base-patch16-256")
encoder = AutoModel.from_pretrained("google/siglip2-base-patch16-256")

# Load moderation head
head = torch.hub.load("khasinski/siglip2-moderation-heads", "nsfw_head")
# Or: head = torch.hub.load("khasinski/siglip2-moderation-heads", "violence_head")
# Or: head = torch.hub.load("khasinski/siglip2-moderation-heads", "ai_generated_head")

# Process image
image = Image.open("test.jpg").convert("RGB")
inputs = processor(images=image, return_tensors="pt")

# Get embedding and classify
with torch.no_grad():
    embedding = encoder.get_image_features(**inputs)
    embedding = embedding / embedding.norm(dim=-1, keepdim=True)  # normalize
    logits = head(embedding)
    probability = torch.sigmoid(logits).item()

print(f"NSFW probability: {probability:.3f}")

Python (ONNX Runtime)

import onnxruntime as ort
import numpy as np
from transformers import AutoModel, AutoProcessor
from PIL import Image

# Load SigLIP2
processor = AutoProcessor.from_pretrained("google/siglip2-base-patch16-256")
encoder = AutoModel.from_pretrained("google/siglip2-base-patch16-256")

# Load ONNX heads
nsfw_session = ort.InferenceSession("nsfw_head.onnx")
violence_session = ort.InferenceSession("violence_head.onnx")
ai_gen_session = ort.InferenceSession("ai_generated_head.onnx")

def get_embedding(image_path):
    image = Image.open(image_path).convert("RGB")
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        embedding = encoder.get_image_features(**inputs)
        embedding = embedding / embedding.norm(dim=-1, keepdim=True)
    return embedding.numpy()

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

# Classify
embedding = get_embedding("test.jpg")

nsfw_prob = sigmoid(nsfw_session.run(None, {"embedding": embedding})[0])[0, 0]
violence_prob = sigmoid(violence_session.run(None, {"embedding": embedding})[0])[0, 0]
ai_gen_prob = sigmoid(ai_gen_session.run(None, {"embedding": embedding})[0])[0, 0]

print(f"NSFW: {nsfw_prob:.3f}, Violence: {violence_prob:.3f}, AI-Gen: {ai_gen_prob:.3f}")

Ruby (ONNX Runtime)

require 'onnxruntime'

# Assumes you have SigLIP2 embeddings from your existing Ruby setup
# (e.g., via the onnxruntime gem with the SigLIP2 ONNX model)

class ModerationClassifier
  def initialize(model_dir = ".")
    @nsfw = OnnxRuntime::Model.new(File.join(model_dir, "nsfw_head.onnx"))
    @violence = OnnxRuntime::Model.new(File.join(model_dir, "violence_head.onnx"))
    @ai_generated = OnnxRuntime::Model.new(File.join(model_dir, "ai_generated_head.onnx"))
  end

  def classify(embedding)
    # embedding should be a normalized 768-dim array
    input = { "embedding" => [embedding] }

    {
      nsfw: sigmoid(@nsfw.predict(input)["logits"][0][0]),
      violence: sigmoid(@violence.predict(input)["logits"][0][0]),
      ai_generated: sigmoid(@ai_generated.predict(input)["logits"][0][0])
    }
  end

  private

  def sigmoid(x)
    1.0 / (1.0 + Math.exp(-x))
  end
end

# Usage
classifier = ModerationClassifier.new("path/to/models")
embedding = get_siglip_embedding("test.jpg")  # your SigLIP2 code
scores = classifier.classify(embedding)

puts "NSFW: #{scores[:nsfw].round(3)}"
puts "Violence: #{scores[:violence].round(3)}"
puts "AI-Generated: #{scores[:ai_generated].round(3)}"

Threshold Recommendations

Training Details

Datasets

Training Configuration

• Optimizer: AdamW (lr=1e-3, weight_decay=0.01)
• Loss: BCEWithLogitsLoss
• Epochs: 15-20
• Batch size: 1024 (on pre-computed embeddings)
• Architecture: Linear(768→256) → ReLU → Dropout(0.2) → Linear(256→1)

Metrics

NSFW Head

              precision    recall  f1-score   support
        Safe       0.96      0.95      0.96      2240
        NSFW       0.97      0.97      0.97      3360
    accuracy                           0.97      5600

Violence Head

              precision    recall  f1-score   support
 Non-violent       0.96      0.97      0.97      1046
     Violent       0.97      0.96      0.97      1169
    accuracy                           0.97      2215

AI-Generated Head

              precision    recall  f1-score   support
        Real       0.92      0.96      0.94      1027
      AI-Gen       0.95      0.90      0.92       835
    accuracy                           0.93      1862

Limitations

• NSFW: Trained primarily on photographic content; may be less accurate on illustrations/anime
• Violence: Trained on real-life violence; graphic illustrations may have lower recall
• AI-Generated: Trained on limited generators (primarily Stable Diffusion variants); newer models like Midjourney v6, DALL-E 3, Flux may evade detection

Files

License

Apache 2.0

Citation

@misc{siglip2-moderation-heads,
  author = {Chris Hasinski},
  title = {SigLIP2 Moderation Heads},
  year = {2025},
  publisher = {HuggingFace},
  url = {https://huggingface.co/khasinski/siglip2-moderation-heads}
}