infer_vit.py

"""
ViT Deepfake Detector - CLI Inference
======================================
Two pre-trained ViT-base-patch16-224 models:
  specialist  -> nano-banana-pro-specialist-h100   (97.5% test accuracy)
  minimal     -> nano-banana-pro-detector-minimal  (95.0% test accuracy)

Both classify: 0 = Real, 1 = AI-Generated

Usage:
  python infer_vit.py image.jpg
  python infer_vit.py image.jpg --model specialist
  python infer_vit.py image.jpg --model minimal
  python infer_vit.py image.jpg --model both       (default, shows ensemble)
  python infer_vit.py image.jpg --threshold 0.6
  python infer_vit.py image.jpg --tta              (test-time augmentation)
  python infer_vit.py image.jpg --gpu              (use GPU, default is CPU)
"""

import sys
import os
import argparse
import warnings
from pathlib import Path

# suppress noisy HF / safetensors / symlink warnings
os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"
warnings.filterwarnings("ignore")

import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from safetensors.torch import load_file

import transformers
transformers.logging.set_verbosity_error()
from transformers import ViTModel, ViTImageProcessor

# ─────────────────────────────────────────────────────────────────────────────
BASE = Path(__file__).parent

MODELS = {
    "specialist": {
        "path":  BASE / "nano-banana-pro-specialist-h100",
        "label": "Specialist  (unfrozen last layer)  97.5% acc",
    },
    "minimal": {
        "path":  BASE / "nano-banana-pro-detector-minimal",
        "label": "Minimal     (frozen backbone)       95.0% acc",
    },
}


# ─────────────────────────────────────────────────────────────────────────────
# ANSI helpers
# ─────────────────────────────────────────────────────────────────────────────
class C:
    RST  = "\033[0m"
    BOLD = "\033[1m"
    DIM  = "\033[2m"
    RED  = "\033[91m"
    GRN  = "\033[92m"
    YLW  = "\033[93m"
    WHT  = "\033[97m"

def pbar(v: float, w: int = 30) -> str:
    n = int(round(v * w))
    return "#" * n + "-" * (w - n)


# ─────────────────────────────────────────────────────────────────────────────
# Model
# ─────────────────────────────────────────────────────────────────────────────
class ViTClassifier(nn.Module):
    """
    ViT-base-patch16-224 backbone (CLS token) + Dropout + Linear(768->2).

    Architecture mirrors the saved checkpoints:
      classifier.0  = Dropout
      classifier.1  = Linear(768, 2)
    """

    def __init__(self, model_dir: Path):
        super().__init__()
        self.vit = ViTModel.from_pretrained(
            str(model_dir), add_pooling_layer=False
        )
        self.classifier = nn.Sequential(
            nn.Dropout(p=0.1),
            nn.Linear(768, 2),
        )

    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
        hidden = self.vit(pixel_values=pixel_values).last_hidden_state
        cls_token = hidden[:, 0]          # (B, 768)
        return self.classifier(cls_token)  # (B, 2)

    @classmethod
    def from_pretrained(cls, model_dir: Path,
                        device: torch.device) -> "ViTClassifier":
        model = cls(model_dir)

        # Load safetensors to CPU first (avoids CUDA init order issues)
        sd = load_file(str(model_dir / "model.safetensors"), device="cpu")

        # Classifier head: safetensors stores "classifier.1.weight/bias"
        # nn.Sequential state_dict uses  "1.weight / 1.bias"
        missing = model.classifier.load_state_dict(
            {"1.weight": sd["classifier.1.weight"],
             "1.bias":   sd["classifier.1.bias"]},
            strict=True,
        )

        return model.to(device).eval()


# ─────────────────────────────────────────────────────────────────────────────
# Preprocessing
# ─────────────────────────────────────────────────────────────────────────────
_processor = None

def get_processor() -> ViTImageProcessor:
    global _processor
    if _processor is None:
        _processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
    return _processor

def to_tensor(img: Image.Image, device: torch.device) -> torch.Tensor:
    return get_processor()(images=img, return_tensors="pt")["pixel_values"].to(device)


# ─────────────────────────────────────────────────────────────────────────────
# Inference
# ─────────────────────────────────────────────────────────────────────────────
import torchvision.transforms.functional as TF

@torch.no_grad()
def run(model: ViTClassifier, img: Image.Image,
        device: torch.device, tta: bool) -> tuple[float, float]:
    """Return (real_prob, ai_prob)."""
    if tta:
        views = [
            img,
            img.transpose(Image.FLIP_LEFT_RIGHT),
            TF.adjust_brightness(img, 1.15),
            TF.adjust_contrast(img, 1.15),
            img.rotate(8),
        ]
        logits = torch.stack([model(to_tensor(v, device)) for v in views]).mean(0)
    else:
        logits = model(to_tensor(img, device))

    probs = F.softmax(logits[0], dim=-1)
    return probs[0].item(), probs[1].item()


# ─────────────────────────────────────────────────────────────────────────────
# Output
# ─────────────────────────────────────────────────────────────────────────────
SEP = "=" * 60

def print_model_result(label: str, real_p: float, ai_p: float, thr: float):
    is_ai   = ai_p >= thr
    verdict = (f"{C.BOLD}{C.RED} AI-GENERATED {C.RST}"
               if is_ai else
               f"{C.BOLD}{C.GRN}    REAL      {C.RST}")

    r_col = C.GRN if real_p > ai_p else C.DIM
    a_col = C.RED if is_ai        else C.DIM

    print(f"\n  {C.BOLD}{label}{C.RST}")
    print(f"  Verdict  : [{verdict}]")
    print(f"  {C.GRN}Real{C.RST}  [{r_col}{pbar(real_p)}{C.RST}] {real_p*100:5.1f}%")
    print(f"  {C.RED}AI  {C.RST}  [{a_col}{pbar(ai_p)}{C.RST}] {ai_p*100:5.1f}%")

def print_ensemble(results: dict):
    real_avg = sum(r for r, _ in results.values()) / len(results)
    ai_avg   = sum(a for _, a in results.values()) / len(results)
    is_ai    = ai_avg >= 0.5
    verdict  = (f"{C.BOLD}{C.RED} AI-GENERATED {C.RST}"
                if is_ai else
                f"{C.BOLD}{C.GRN}    REAL      {C.RST}")

    r_col = C.GRN if real_avg > ai_avg else C.DIM
    a_col = C.RED if is_ai             else C.DIM

    print(f"\n  {C.BOLD}{C.WHT}Ensemble (average of both models){C.RST}")
    print(f"  Verdict  : [{verdict}]")
    print(f"  {C.GRN}Real{C.RST}  [{r_col}{pbar(real_avg)}{C.RST}] {real_avg*100:5.1f}%")
    print(f"  {C.RED}AI  {C.RST}  [{a_col}{pbar(ai_avg)}{C.RST}] {ai_avg*100:5.1f}%")


# ─────────────────────────────────────────────────────────────────────────────
# CLI
# ─────────────────────────────────────────────────────────────────────────────
def main():
    parser = argparse.ArgumentParser(
        description="ViT deepfake/AI-image detection",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument("image",
                        help="Path to image file (jpg/png/webp/bmp)")
    parser.add_argument("--model", "-m",
                        choices=["specialist", "minimal", "both"],
                        default="both",
                        help="Model to use (default: both + ensemble)")
    parser.add_argument("--threshold", "-t",
                        type=float, default=0.5,
                        help="AI-generated probability threshold (default: 0.5)")
    parser.add_argument("--tta",
                        action="store_true",
                        help="Test-time augmentation: average 5 views")
    parser.add_argument("--gpu",
                        action="store_true",
                        help="Use GPU (default: CPU, fast enough for ViT-base)")
    args = parser.parse_args()

    # ── validate ──────────────────────────────────────────────────────
    img_path = Path(args.image)
    if not img_path.exists():
        sys.exit(f"{C.RED}Error: file not found: {img_path}{C.RST}")

    try:
        img = Image.open(img_path).convert("RGB")
    except Exception as e:
        sys.exit(f"{C.RED}Error: cannot open image: {e}{C.RST}")

    # Default to CPU: ViT-base (~330 MB) is fast enough on CPU for single images,
    # and avoids conflicts with other GPU processes (e.g. ManipShield training).
    # Pass --gpu to force GPU if you have free VRAM.
    use_gpu = getattr(args, "gpu", False)
    device = torch.device(
        "cuda" if use_gpu and torch.cuda.is_available() else "cpu"
    )
    models_to_run = (["specialist", "minimal"]
                     if args.model == "both" else [args.model])

    # ── header ────────────────────────────────────────────────────────
    print(f"\n{C.BOLD}{C.WHT}{SEP}{C.RST}")
    print(f"{C.BOLD}{C.WHT}  ViT Deepfake / AI-Image Detection{C.RST}")
    print(f"{C.BOLD}{C.WHT}{SEP}{C.RST}")
    print(f"  Image     : {img_path}")
    print(f"  Size      : {img.width} x {img.height}")
    print(f"  Device    : {device}")
    print(f"  Threshold : {args.threshold}")
    print(f"  TTA       : {'yes (5 views)' if args.tta else 'no'}")

    # ── preload processor ─────────────────────────────────────────────
    print(f"\n{C.DIM}  Loading ViT processor...{C.RST}", end=" ", flush=True)
    get_processor()
    print(f"{C.DIM}done{C.RST}")

    # ── load all models first, then infer ────────────────────────────
    loaded = {}
    for key in models_to_run:
        print(f"{C.DIM}  Loading {key} weights...{C.RST}", end=" ", flush=True)
        loaded[key] = ViTClassifier.from_pretrained(MODELS[key]["path"], device)
        print(f"{C.DIM}done{C.RST}")

    results = {}
    for key, model in loaded.items():
        real_p, ai_p = run(model, img, device, args.tta)
        results[key] = (real_p, ai_p)
        print_model_result(MODELS[key]["label"], real_p, ai_p, args.threshold)

    if len(results) == 2:
        print_ensemble(results)

    print(f"\n{C.BOLD}{C.WHT}{SEP}{C.RST}\n")


if __name__ == "__main__":
    main()