download_model.py

"""
DeepGuard — Model Download Script
Downloads the ViT-based deepfake detection ONNX model from Hugging Face Hub.

Usage:
    python download_model.py

Model: onnx-community/Deep-Fake-Detector-v2-Model-ONNX
  - Architecture: google/vit-base-patch16-224-in21k (fine-tuned)
  - Task:         Binary classification — Realism vs. Deepfake
  - Labels:       {0: "Realism", 1: "Deepfake"}
"""

import os
import sys
import hashlib

MODELS_DIR = os.path.join(os.path.dirname(__file__), "models")
MODEL_DEST = os.path.join(MODELS_DIR, "deepfake_vit.onnx")

REPO_ID = "onnx-community/Deep-Fake-Detector-v2-Model-ONNX"
FILENAME = "onnx/model.onnx"  # Path inside the HF repo


def sha256(path: str) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(65536), b""):
            h.update(chunk)
    return h.hexdigest()


def download():
    os.makedirs(MODELS_DIR, exist_ok=True)

    if os.path.exists(MODEL_DEST):
        size_mb = os.path.getsize(MODEL_DEST) / (1024 * 1024)
        print(f"[DeepGuard] Model already exists ({size_mb:.1f} MB): {MODEL_DEST}")
        print(f"[DeepGuard] SHA-256: {sha256(MODEL_DEST)}")
        print("[DeepGuard] Delete the file and re-run this script to force re-download.")
        return

    print(f"[DeepGuard] Downloading model from Hugging Face Hub...")
    print(f"  Repo:    {REPO_ID}")
    print(f"  File:    {FILENAME}")
    print(f"  Target:  {MODEL_DEST}")
    print()

    try:
        from huggingface_hub import hf_hub_download

        tmp_path = hf_hub_download(
            repo_id=REPO_ID,
            filename=FILENAME,
            cache_dir=MODELS_DIR,
            local_dir=MODELS_DIR,
            local_dir_use_symlinks=False,
        )

        # hf_hub_download writes to local_dir/<filename>
        # Move if needed
        expected_local = os.path.join(MODELS_DIR, "onnx", "model.onnx")
        if os.path.exists(expected_local) and not os.path.exists(MODEL_DEST):
            import shutil
            shutil.move(expected_local, MODEL_DEST)

        if not os.path.exists(MODEL_DEST):
            # Try symlink/copy from tmp_path
            import shutil
            shutil.copy2(tmp_path, MODEL_DEST)

        size_mb = os.path.getsize(MODEL_DEST) / (1024 * 1024)
        checksum = sha256(MODEL_DEST)
        print(f"\n[DeepGuard] [OK] Download complete!")
        print(f"  Size:    {size_mb:.1f} MB")
        print(f"  SHA-256: {checksum}")
        print(f"  Path:    {MODEL_DEST}")

    except ImportError:
        print("[ERROR] huggingface_hub not installed. Run: pip install huggingface_hub")
        sys.exit(1)
    except Exception as e:
        print(f"[ERROR] Download failed: {e}")
        print()
        print("Manual download instructions:")
        print(f"  1. Visit: https://huggingface.co/{REPO_ID}/tree/main/onnx")
        print(f"  2. Download 'model.onnx'")
        print(f"  3. Place it at: {MODEL_DEST}")
        sys.exit(1)


if __name__ == "__main__":
    download()