ViT Base: Export artifacts + push to Hugging Face Hub (fusion-ready)

[lukhsaankumar]

Issue Type

• Model: ML model bug, training issue, or architecture problem
• Data: Dataset issue, preprocessing bug, or data pipeline problem
• Web: Frontend bug or UI issue in the Next.js dashboard
• API: Backend API bug or FastAPI endpoint issue
• Research: Research question or experimental feature request
• Documentation: Documentation bug or improvement needed
• Bug: General bug fix needed
• Enhancement: New feature or improvement request

Description

Standardize the Vision Transformer (ViT Base) model export workflow and publish all required artifacts to Hugging Face Hub so the fusion model and backend can load it remotely in a reproducible way.

This ticket updates the workflow from:

• Previous: save artifacts locally under models/vit_base/
• New: export artifacts → upload to Hugging Face → backend downloads + loads for inference

Hugging Face repo: DeepFakeDetector/vit-base (https://huggingface.co/DeepFakeDetector/vit-base)

Acceptance Criteria

• A Hugging Face Write token is created and available locally as an environment variable (HF_TOKEN_WRITE)
• vit_transfer_baseline.ipynb contains a final Export + Upload section
• Artifacts are exported with the standardized layout:
  • model.pth (or HF-compatible weights)
  • config.json (architecture, patch size, hidden dim, threshold, notes)
  • preprocess.json (resize, normalization, interpolation)
  • label_map.json (0 = real, 1 = fake)
  • README.md (model card)
• All artifacts are uploaded to the Hugging Face repo
• A simple reload sanity check (download → load → eval) is included in the notebook

Deliverables

Hugging Face (canonical storage)

Uploaded to: https://huggingface.co/DeepFakeDetector/vit-base

Must contain:

• model.pth
• config.json
• preprocess.json
• label_map.json
• README.md

Notebook update

vit_transfer_baseline.ipynb includes cells that:

1. Export all artifacts
2. Upload them to Hugging Face

Additional Context

Label Convention (CRITICAL)

Label convention is fixed and shared across all models:

• 0 = real
• 1 = fake

Preprocessing

• Preprocessing details must match training exactly (ViT is sensitive to this)

Output

• Output is binary classification; fusion will consume prob_fake
• Embeddings are optional but encouraged if already available from the ViT head

---

Implementation Notes (rough idea)

1) Create Hugging Face Write Token

Create a write token in Hugging Face account settings and export locally:

Environment variable (recommended):

export HF_TOKEN_WRITE=hf_xxxxxxxxxxxxxxxxx

Never commit this token to git.

2) Notebook Cell Snippets to Add

Add these cells to the end of vit_transfer_baseline.ipynb:

Cell A — Export ViT artifacts

import os, json, torch

EXPORT_DIR = "export/vit-base"
os.makedirs(EXPORT_DIR, exist_ok=True)

# Save model weights (state_dict)
torch.save(model.state_dict(), os.path.join(EXPORT_DIR, "model.pth"))

# Preprocessing details (must match training!)
preprocess = {
    "input_size": 224,
    "resize": 256,
    "center_crop": True,
    "interpolation": "bicubic",
    "normalize": {
        "mean": [0.5, 0.5, 0.5],
        "std":  [0.5, 0.5, 0.5]
    }
}

label_map = {"0": "real", "1": "fake"}

config = {
    "name": "vit-base",
    "framework": "pytorch",
    "arch": "vit_base_patch16_224",
    "patch_size": 16,
    "hidden_dim": 768,
    "num_layers": 12,
    "num_heads": 12,
    "num_classes": 2,
    "threshold": 0.50,
    "labels": label_map,
    "notes": "ViT Base transfer learning model for deepfake detection"
}

with open(os.path.join(EXPORT_DIR, "preprocess.json"), "w") as f:
    json.dump(preprocess, f, indent=2)

with open(os.path.join(EXPORT_DIR, "label_map.json"), "w") as f:
    json.dump(label_map, f, indent=2)

with open(os.path.join(EXPORT_DIR, "config.json"), "w") as f:
    json.dump(config, f, indent=2)

readme = """---
license: apache-2.0
tags:
  - vision-transformer
  - deepfake-detection
  - image-classification
---

# ViT Base – DeepFakeDetector

Vision Transformer (ViT-B/16) fine-tuned for binary deepfake detection.

## Labels
- 0 = real
- 1 = fake

## Output
- prob_fake ∈ [0,1] (softmax or sigmoid)
"""

with open(os.path.join(EXPORT_DIR, "README.md"), "w") as f:
    f.write(readme)

print("Exported ViT Base artifacts to:", EXPORT_DIR)

Cell B — Upload artifacts to Hugging Face Hub

import os
from huggingface_hub import HfApi

REPO_ID = "DeepFakeDetector/vit-base"
HF_TOKEN = os.environ.get("HF_TOKEN_WRITE")

assert HF_TOKEN, "HF_TOKEN_WRITE env var not set."

api = HfApi(token=HF_TOKEN)

# Create repo if it doesn't exist
api.create_repo(repo_id=REPO_ID, repo_type="model", exist_ok=True)

api.upload_folder(
    folder_path=EXPORT_DIR,
    repo_id=REPO_ID,
    repo_type="model",
    commit_message="Upload ViT Base artifacts (weights + configs + preprocess)"
)

print(f"Uploaded to https://huggingface.co/{REPO_ID}")

Cell C — (Optional but recommended) Sanity reload test

import torch
from huggingface_hub import snapshot_download

local_repo = snapshot_download(repo_id=REPO_ID, repo_type="model")

state = torch.load(os.path.join(local_repo, "model.pth"), map_location="cpu")
model.load_state_dict(state)
model.eval()

print("ViT Base weights reloaded successfully from HF")

Definition of Done

• Hugging Face repo contains all required artifacts
• Notebook can export + upload in one run
• Reload test passes
• Preprocessing + label convention documented and consistent
• Model is ready for fusion + backend consumption