""" model.py -------- Loads the pretrained deepfake detection model. WHERE THE WEIGHTS COME FROM --------------------------- This project uses the real, publicly available Hugging Face model: https://huggingface.co/dima806/deepfake_vs_real_image_detection It is a Vision Transformer (google/vit-base-patch16-224-in21k) fine-tuned by the community author "dima806" on ~2.8k real and deepfake face images for the binary task REAL vs FAKE. It is free to use for academic purposes (the base ViT is Apache-2.0 licensed) and is small enough (~350 MB fp32) for CPU inference. On first run this module downloads the weights from the Hugging Face Hub, converts them into a single local file (models/pretrained_model.pth) that contains the state dict + model config + preprocessing parameters, and every later run loads 100% locally from disk (no internet needed). You can also pre-download it manually: python download_model.py Environment variables: DEEPFAKE_MODEL_ID any compatible Hugging Face image-classification model (default: dima806/deepfake_vs_real_image_detection) DEEPFAKE_MODEL_DIR folder where pretrained_model.pth is stored DEEPFAKE_QUANTIZE "1" = apply dynamic int8 quantization after loading (smaller RAM footprint, slightly slower/less accurate) """ import os import threading import time import numpy as np import torch from PIL import Image from torch import nn from transformers import ViTConfig, ViTForImageClassification BASE_DIR = os.path.dirname(os.path.abspath(__file__)) MODELS_DIR = os.environ.get("DEEPFAKE_MODEL_DIR", os.path.join(BASE_DIR, "models")) MODEL_PATH = os.path.join(MODELS_DIR, "pretrained_model.pth") HF_MODEL_ID = os.environ.get("DEEPFAKE_MODEL_ID", "dima806/deepfake_vs_real_image_detection") QUANTIZE = os.environ.get("DEEPFAKE_QUANTIZE", "0") == "1" torch.set_num_threads(max(1, min(4, os.cpu_count() or 2))) class DeepfakeModel: """Singleton wrapper around the pretrained ViT deepfake classifier.""" def __init__(self): self._model = None self._lock = threading.Lock() self.size = 224 self.mean = [0.5, 0.5, 0.5] self.std = [0.5, 0.5, 0.5] self.labels = {0: "REAL", 1: "FAKE"} @property def is_loaded(self): return self._model is not None # ------------------------------------------------------------------ # # Loading # # ------------------------------------------------------------------ # def load(self): with self._lock: if self._model is not None: return if not os.path.exists(MODEL_PATH): print(f"[model] {MODEL_PATH} not found -> downloading " f"'{HF_MODEL_ID}' from the Hugging Face Hub ...") self._download() print(f"[model] loading local weights: {MODEL_PATH}") t0 = time.time() try: # mmap keeps weight tensors on disk until touched, which # halves the peak RAM during loading (matters on 512 MB hosts) bundle = torch.load( MODEL_PATH, map_location="cpu", weights_only=True, mmap=True ) except TypeError: # older torch without mmap support bundle = torch.load(MODEL_PATH, map_location="cpu", weights_only=True) config = ViTConfig.from_dict(bundle["config"]) net = ViTForImageClassification(config) net.load_state_dict(bundle["state_dict"], strict=True) self.size = int(bundle.get("image_size", 224)) self.mean = list(bundle.get("image_mean", self.mean)) self.std = list(bundle.get("image_std", self.std)) id2label = getattr(config, "id2label", None) or {} if id2label: self.labels = {int(k): str(v).upper() for k, v in id2label.items()} if QUANTIZE: try: net = torch.quantization.quantize_dynamic( net, {nn.Linear}, dtype=torch.qint8 ) print("[model] dynamic int8 quantization applied") except Exception as exc: # pragma: no cover - optional path print(f"[model] quantization skipped ({exc})") net.eval() self._model = net self._warmup() print(f"[model] ready in {time.time() - t0:.1f}s " f"(labels: {self.labels})") def _download(self): """Download from the Hugging Face Hub and save a single local .pth.""" from transformers import AutoModelForImageClassification, AutoImageProcessor os.makedirs(MODELS_DIR, exist_ok=True) processor = AutoImageProcessor.from_pretrained(HF_MODEL_ID) net = AutoModelForImageClassification.from_pretrained(HF_MODEL_ID) net.eval() raw_size = processor.size if hasattr(raw_size, "get"): size = int(raw_size.get( "shortest_edge", raw_size.get("height", raw_size.get("width", 224)) )) else: size = int(getattr( raw_size, "shortest_edge", getattr(raw_size, "height", getattr(raw_size, "width", 224)), )) bundle = { "format_version": 1, "source_model_id": HF_MODEL_ID, "config": net.config.to_dict(), "state_dict": net.state_dict(), "image_size": size, "image_mean": list(processor.image_mean), "image_std": list(processor.image_std), } tmp_path = MODEL_PATH + ".tmp" torch.save(bundle, tmp_path) os.replace(tmp_path, MODEL_PATH) print(f"[model] saved {os.path.getsize(MODEL_PATH) / 1e6:.0f} MB -> {MODEL_PATH}") def _warmup(self): dummy = Image.new("RGB", (self.size, self.size), (128, 128, 128)) try: with torch.no_grad(): self._forward(dummy) except Exception as exc: # pragma: no cover print(f"[model] warmup warning: {exc}") # ------------------------------------------------------------------ # # Inference # # ------------------------------------------------------------------ # def preprocess(self, pil_img): img = pil_img.convert("RGB").resize((self.size, self.size), Image.BICUBIC) arr = np.asarray(img, dtype=np.float32) / 255.0 arr = (arr - np.asarray(self.mean, dtype=np.float32)) / \ np.asarray(self.std, dtype=np.float32) tensor = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0) return tensor def _forward(self, pil_img): inputs = self.preprocess(pil_img) outputs = self._model(pixel_values=inputs) probs = torch.softmax(outputs.logits[0], dim=0) return probs def predict(self, pil_img): """ Returns (label_uppercase, confidence_percent, prob_fake_0to1). Label is normalized to 'REAL' or 'FAKE'. """ self.load() probs = self._forward(pil_img) idx = int(torch.argmax(probs).item()) label_raw = self.labels.get(idx, f"CLASS_{idx}").upper() fake_idx = next((i for i, l in self.labels.items() if "FAKE" in l or l.startswith("DF")), 1) real_idx = next((i for i, l in self.labels.items() if "REAL" in l or "GENUINE" in l), 0) prob_fake = float(probs[fake_idx].item()) prob_real = float(probs[real_idx].item()) if "FAKE" in label_raw: label = "FAKE" confidence = prob_fake * 100.0 else: label = "REAL" confidence = prob_real * 100.0 return label, round(confidence, 2), prob_fake _model_singleton = DeepfakeModel() def get_model(): _model_singleton.load() return _model_singleton def ensure_model_downloaded(): get_model() return MODEL_PATH if __name__ == "__main__": path = ensure_model_downloaded() print("Model available at:", path)