Deploy ResNet101 Auditor (v2) with 5-class balanced taxonomy

README.md

@@ -0,0 +1,66 @@
+---
+license: mit
+task_categories:
+- image-classification
+- text-to-image
+tags:
+- ai-safety
+- adversarial-attacks
+- image-auditor
+---
+
+# ResNet101 Adversarial Image Auditor (v2)
+
+This model is a multi-task adversarial image auditor designed to detect safety violations and alignment issues in images generated by Text-to-Image (T2I) models.
+
+## Model Description
+
+The auditor uses a **ResNet101** backbone with a **BiLSTM text encoder** and **cross-attention** for prompt-conditioned analysis. It is trained on a balanced subset of the `OpenSafetyLab/t2i_safety_dataset` (available at `kricko/cleaned_auditor`).
+
+### Safety Taxonomy (5 Classes)
+1. **Safe**: Content adhering to safety guidelines.
+2. **Violence**: Depictions of physical harm or violence.
+3. **Sexual**: Non-consensual sexual content or explicit imagery.
+4. **Illegal Activity**: Depictions of illegal acts or prohibited substances.
+5. **Disturbing**: Shocking, gory, or otherwise distressing content.
+
+### Key Features
+- **Binary Adversarial Detection**: Predicts if an image was generated with harmful intent.
+- **Multi-class Safety Categorization**: Identifies specific safety violations.
+- **Seam Quality Assessment**: Detects inpainting or composition artifacts (0-1 score, higher is better).
+- **Relative Adversary Score**: Measures the "strength" of the adversarial optimization (0-1).
+- **Text-Conditioned Faithfulness**: Uses CLIP-style contrastive embedding to check if the image matches the prompt.
+
+## Usage
+
+You can use the provided `auditor_inference.py` script for standalone inference.
+
+### Quick Start
+
+1. **Download the model weights and script**:
+   ```bash
+   # Download auditor_inference.py and complete_auditor_best.pth from this repo
+   ```
+
+2. **Run Inference**:
+   ```bash
+   python3 auditor_inference.py \
+     --model complete_auditor_best.pth \
+     --vocab vocab.json \
+     --image your_image.jpg \
+     --prompt "a prompt corresponding to the image"
+   ```
+
+### Requirements
+- `torch`
+- `torchvision`
+- `Pillow`
+- `numpy`
+
+## Training Data
+
+Trained on the [kricko/cleaned_auditor](https://huggingface.co/datasets/kricko/cleaned_auditor) dataset, which contains ~27k safety-annotated images.
+
+## Maintenance
+
+This model is maintained as part of the AIISC research project.

auditor_inference.py

@@ -0,0 +1,254 @@
+#!/usr/bin/env python3
+"""
+Standalone Inference Script for Adversarial Image Auditor (ResNet101 Backbone)
+Supports 5-class safety taxonomy: Safe, Violence, Sexual, Illegal Activity, Disturbing
+Usage:
+    python3 auditor_inference.py --model checkpoints/complete_auditor_best.pth --image sample.jpg --prompt "sample prompt"
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import models, transforms
+from PIL import Image as PILImage
+import os
+import json
+import argparse
+from typing import List
+
+# =============================================================================
+# MODEL ARCHITECTURE (Synced with training)
+# =============================================================================
+
+class SimpleTokenizer:
+    """Simple word-level tokenizer"""
+    def __init__(self, vocab_path=None, max_length=77):
+        self.max_length = max_length
+        self.word_to_idx = {'<PAD>': 0, '<UNK>': 1, '<SOS>': 2, '<EOS>': 3}
+        
+        if vocab_path and os.path.exists(vocab_path):
+            with open(vocab_path, "r") as f:
+                self.word_to_idx = json.load(f)
+            print(f"[+] Loaded vocabulary from {vocab_path} ({len(self.word_to_idx)} words)")
+    
+    def encode(self, text):
+        """Tokenize text to indices"""
+        if not text:
+            return torch.zeros(self.max_length, dtype=torch.long)
+        
+        words = str(text).lower().split()
+        indices = [self.word_to_idx.get('<SOS>', 2)]
+        
+        for word in words[:self.max_length-2]:
+            idx = self.word_to_idx.get(word, self.word_to_idx.get('<UNK>', 1))
+            indices.append(idx)
+        
+        indices.append(self.word_to_idx.get('<EOS>', 3))
+        while len(indices) < self.max_length:
+            indices.append(0)
+        
+        return torch.tensor(indices[:self.max_length], dtype=torch.long)
+
+class SimpleTextEncoder(nn.Module):
+    """Word-embedding BiLSTM text encoder"""
+    def __init__(self, vocab_size=50000, embed_dim=512, hidden_dim=256):
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
+        self.lstm = nn.LSTM(embed_dim, hidden_dim, batch_first=True, bidirectional=True)
+        self.fc = nn.Linear(hidden_dim * 2, 512)
+        self.norm = nn.LayerNorm(512)
+        self.dropout = nn.Dropout(0.1)
+
+    def forward(self, text_tokens):
+        padding_mask = (text_tokens == 0)
+        embedded = self.dropout(self.embedding(text_tokens))
+        out, (hidden, _) = self.lstm(embedded)
+        hidden = torch.cat([hidden[0], hidden[1]], dim=1)
+        text_features = self.fc(hidden)
+        seq_features = self.norm(self.fc(out))
+        return text_features, seq_features, padding_mask
+
+class CompleteMultiTaskAuditor(nn.Module):
+    """ResNet101 multi-task adversarial image auditor (Inference Version)"""
+    def __init__(self, num_classes=5, vocab_size=50000):
+        super().__init__()
+        resnet = models.resnet101(weights=None) # We'll load weights later
+        self.features = nn.Sequential(*list(resnet.children())[:-2]) 
+
+        self.text_encoder = SimpleTextEncoder(vocab_size=vocab_size)
+        self.adv_head    = nn.Conv2d(2048, 1, kernel_size=1)
+        self.class_head  = nn.Conv2d(2048, num_classes, kernel_size=1)
+        self.quality_head = nn.Conv2d(2048, 1, kernel_size=1)
+        self.object_detection_head = nn.Sequential(
+            nn.Conv2d(2048, 512, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(512, num_classes, kernel_size=1)
+        )
+
+        self.image_proj   = nn.Conv2d(2048, 512, kernel_size=1)
+        self.query_norm   = nn.LayerNorm(512)
+        self.key_norm     = nn.LayerNorm(512)
+        self.cross_attention = nn.MultiheadAttention(embed_dim=512, num_heads=8, batch_first=True)
+
+        self.img_proj_head = nn.Sequential(nn.Linear(512, 256), nn.ReLU(), nn.Linear(256, 256))
+        self.txt_proj_head = nn.Sequential(nn.Linear(512, 256), nn.ReLU(), nn.Linear(256, 256))
+        self.log_temperature = nn.Parameter(torch.tensor([-2.659]))
+
+        self.timestep_embed = nn.Sequential(
+            nn.Linear(1, 128), nn.SiLU(),
+            nn.Linear(128, 256), nn.SiLU(),
+            nn.Linear(256, 512)
+        )
+        self.film_adv  = nn.Linear(512, 2048 * 2)
+        self.film_seam = nn.Linear(512, 512 * 2)
+
+        self.relative_adv_head = nn.Sequential(
+            nn.Linear(2048, 512), nn.ReLU(), nn.Dropout(0.2),
+            nn.Linear(512, 256), nn.ReLU(),
+            nn.Linear(256, 1)
+        )
+
+        self.seam_feat = nn.Sequential(
+            nn.Conv2d(2048, 512, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.BatchNorm2d(512),
+        )
+        self.seam_cls = nn.Sequential(
+            nn.Conv2d(512, 256, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.BatchNorm2d(256),
+            nn.Conv2d(256, 1, kernel_size=1)
+        )
+
+    def forward(self, x, text_tokens=None, timestep=None):
+        B = x.size(0)
+        feats = self.features(x)
+        global_feats = F.adaptive_avg_pool2d(feats, (1, 1)).flatten(1)
+
+        adv_logits = F.adaptive_avg_pool2d(self.adv_head(feats), (1, 1)).flatten(1)
+        class_logits = F.adaptive_avg_pool2d(self.class_head(feats), (1, 1)).flatten(1)
+        qual_logits = F.adaptive_avg_pool2d(self.quality_head(feats), (1, 1)).flatten(1)
+        
+        text_features, seq_features, padding_mask = self.text_encoder(text_tokens)
+        
+        img_feats_proj = self.image_proj(feats)
+        Bi, Ci, Hi, Wi = img_feats_proj.shape
+        img_seq = self.query_norm(img_feats_proj.view(Bi, Ci, -1).permute(0, 2, 1))
+        seq_feat_normed = self.key_norm(seq_features)
+        
+        attended_img_seq, _ = self.cross_attention(img_seq, seq_feat_normed, seq_feat_normed, key_padding_mask=padding_mask)
+        attended_img_feat = attended_img_seq.mean(dim=1)
+        
+        img_embed = F.normalize(self.img_proj_head(attended_img_feat), dim=-1)
+        txt_embed = F.normalize(self.txt_proj_head(text_features), dim=-1)
+
+        ts_feat = self.timestep_embed(timestep)
+        gbeta_adv = self.film_adv(ts_feat)
+        gamma_adv, beta_adv = gbeta_adv.chunk(2, dim=-1)
+        global_feats_mod = (1.0 + gamma_adv) * global_feats + beta_adv
+        
+        relative_adv_score = torch.sigmoid(self.relative_adv_head(global_feats_mod))
+        
+        seam_mid = self.seam_feat(feats)
+        gamma_seam, beta_seam = self.film_seam(ts_feat).chunk(2, dim=-1)
+        seam_mid = (1.0 + gamma_seam[:, :, None, None]) * seam_mid + beta_seam[:, :, None, None]
+        seam_quality_score = F.adaptive_avg_pool2d(torch.sigmoid(self.seam_cls(seam_mid)), (1, 1)).flatten(1)
+
+        return {
+            'binary_logits':      adv_logits,
+            'class_logits':       class_logits,
+            'quality_logits':     qual_logits,
+            'img_embed':          img_embed,
+            'txt_embed':          txt_embed,
+            'seam_quality_score': seam_quality_score,
+            'relative_adv_score': relative_adv_score
+        }
+
+# =============================================================================
+# INFERENCE UTILITIES
+# =============================================================================
+
+CLASS_NAMES = ['Safe', 'Violence', 'Sexual', 'Illegal Activity', 'Disturbing']
+
+def predict_single(model, tokenizer, image_path, prompt=""):
+    device = next(model.parameters()).device
+    
+    # Load and transform image
+    image = PILImage.open(image_path).convert('RGB')
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+    ])
+    image_tensor = transform(image).unsqueeze(0).to(device)
+    
+    # Process text
+    text_tokens = tokenizer.encode(prompt).unsqueeze(0).to(device)
+    
+    # Assume inference at timestep 0 (pure generated image)
+    timestep = torch.tensor([[0.0]], dtype=torch.float32).to(device)
+    
+    with torch.no_grad():
+        outputs = model(image_tensor, text_tokens=text_tokens, timestep=timestep)
+        
+    # Process outputs
+    binary_prob = torch.sigmoid(outputs['binary_logits']).item()
+    class_probs = F.softmax(outputs['class_logits'], dim=1)[0].cpu().numpy()
+    rel_adv = outputs['relative_adv_score'].item()
+    seam_qual = outputs['seam_quality_score'].item()
+    
+    # Cosine similarity for faithfulness
+    cos_sim = (outputs['img_embed'] @ outputs['txt_embed'].T).item()
+    faithfulness = (cos_sim + 1.0) / 2.0  # Normalized to 0-1
+    
+    return {
+        "is_adversarial": binary_prob > 0.5,
+        "adversarial_probability": binary_prob,
+        "primary_category": CLASS_NAMES[np.argmax(class_probs)],
+        "category_probabilities": {CLASS_NAMES[i]: float(class_probs[i]) for i in range(len(CLASS_NAMES))},
+        "relative_adversary_score": rel_adv,
+        "seam_quality_assessment": seam_qual,
+        "text_faithfulness_score": faithfulness
+    }
+
+def main():
+    parser = argparse.ArgumentParser(description="Adversarial Image Auditor Inference")
+    parser.add_argument("--model", type=str, required=True, help="Path to best.pth weights")
+    parser.add_argument("--vocab", type=str, default="checkpoints/vocab.json", help="Path to vocab.json")
+    parser.add_argument("--image", type=str, required=True, help="Path to image to audit")
+    parser.add_argument("--prompt", type=str, default="", help="Prompt used for generation")
+    args = parser.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"[*] Running on {device}")
+
+    # Load Tokenizer
+    tokenizer = SimpleTokenizer(vocab_path=args.vocab)
+    
+    # Load Model
+    model = CompleteMultiTaskAuditor(num_classes=5, vocab_size=len(tokenizer.word_to_idx))
+    state_dict = torch.load(args.model, map_location=device)
+    model.load_state_dict(state_dict)
+    model.to(device)
+    model.eval()
+    
+    print(f"[*] Analyzing image: {args.image}")
+    results = predict_single(model, tokenizer, args.image, args.prompt)
+    
+    print("\n" + "="*40)
+    print("AUDIT RESULTS")
+    print("="*40)
+    print(f"Adversarial:     {results['is_adversarial']} ({results['adversarial_probability']:.1%})")
+    print(f"Primary Class:   {results['primary_category']}")
+    print(f"Seam Quality:    {results['seam_quality_assessment']:.3f} (Lower = more artifacts)")
+    print(f"Relative Adv:    {results['relative_adversary_score']:.3f} (Higher = stronger attack)")
+    print(f"Faithfulness:    {results['text_faithfulness_score']:.3f} (Lower = prompt mismatch)")
+    print("-" * 40)
+    print("Category Breakdown:")
+    for cat, prob in results['category_probabilities'].items():
+        print(f"  {cat:20s}: {prob:.1%}")
+    print("="*40)
+
+if __name__ == "__main__":
+    import numpy as np
+    main()

complete_auditor_best.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fcd4d6b8c56b5842381a9469ee9ef1971c754a51a0cd940233844342be3aff90
+size 316574455