Spaces:
Sleeping
Sleeping
OmPranithSai
Configure Render free tier: quantization, build-time model download, mmap loading
98dcd01 | """ | |
| 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"} | |
| 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) | |