Upload pixelprism_detector.py with huggingface_hub

pixelprism_detector.py

@@ -0,0 +1,98 @@
+"""PixelPrism v0.1 — DeepfakeDetector subclass for the BitMind DFD Arena.
+
+This is a sanitized single-detector submission representing the most
+informative component of PixelPrism's production V9 16-detector ensemble:
+the Swin V2 transformer head (haywoodsloan/ai-image-detector-deploy),
+which V9's permutation-importance audit ranked at 0.271 (38% of total
+discriminative power, the strongest single feature in our ensemble).
+
+The full PixelPrism V9 ensemble fuses 16 detectors via a
+HistGradientBoostingClassifier meta-classifier and runs against a live
+production API at https://pixelprism.ai. Some V9 components depend on
+non-MIT-redistributable weights (FLUX-schnell for DIRE-FLUX, SDXL for
+DIRE-SDXL) so this submission is the open-redistributable subset.
+
+For the full 16-detector live numbers see https://pixelprism.ai/leaderboard
+(refreshed monthly with each retrain).
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import torch
+from huggingface_hub import hf_hub_download
+from PIL import Image
+from transformers import AutoImageProcessor, Swinv2ForImageClassification
+
+from arena.detectors.deepfake_detectors import DeepfakeDetector
+from arena.detectors import DETECTOR_REGISTRY
+
+
+# Module-level cache so repeated instantiations don't reload the model.
+_PROCESSOR = None
+_MODEL = None
+
+
+@DETECTOR_REGISTRY.register_module(module_name='PixelPrism')
+class PixelPrismDetector(DeepfakeDetector):
+    """Single-detector submission: Swin V2 wrapped with PixelPrism's production
+    inference path. Outputs P(AI) ∈ [0, 1] per image.
+
+    Attributes set from the YAML config:
+        hf_repo (str):     "pixelprism-ai/dfd-arena-mini"
+        backbone_repo (str): "haywoodsloan/ai-image-detector-deploy"
+        ai_label_idx (int):  the class index that means "AI" in the Swin head (0)
+    """
+
+    def __init__(
+        self,
+        model_name: str = 'PixelPrism',
+        config: str = 'pixelprism_config.yaml',
+        cuda: bool = True,
+    ):
+        super().__init__(model_name, config, cuda)
+
+    # ───────────────────────── model load ─────────────────────────
+
+    def load_model(self):
+        global _PROCESSOR, _MODEL
+
+        if _MODEL is not None and _PROCESSOR is not None:
+            self.processor = _PROCESSOR
+            self.model = _MODEL
+            return
+
+        # The Swin V2 model + preprocessor come from haywoodsloan's MIT-licensed
+        # repo; we re-host the same files at pixelprism-ai/dfd-arena-mini so the
+        # arena evaluator can pull everything from one place.
+        repo = self.hf_repo
+
+        config_path = hf_hub_download(repo, 'config.json')
+        weights_path = hf_hub_download(repo, 'model.safetensors')
+        preproc_path = hf_hub_download(repo, 'preprocessor_config.json')
+
+        self.processor = AutoImageProcessor.from_pretrained(Path(preproc_path).parent)
+        self.model = Swinv2ForImageClassification.from_pretrained(Path(weights_path).parent)
+        self.model.to(self.device).eval()
+
+        _PROCESSOR = self.processor
+        _MODEL = self.model
+
+    # ───────────────────────── inference ─────────────────────────
+
+    def preprocess(self, image: Image.Image) -> torch.Tensor:
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        inputs = self.processor(images=image, return_tensors='pt')
+        return inputs['pixel_values'].to(self.device)
+
+    def __call__(self, image: Image.Image) -> float:
+        """Returns P(AI) ∈ [0, 1] for a single PIL image."""
+        pixel_values = self.preprocess(image)
+        with torch.no_grad():
+            logits = self.model(pixel_values=pixel_values).logits
+            probs = torch.softmax(logits, dim=-1).cpu().numpy().ravel()
+        # ai_label_idx is set from the YAML config (typically 0 for haywoodsloan)
+        ai_idx = int(getattr(self, 'ai_label_idx', 0))
+        return float(probs[ai_idx])