DeiT (Distilled): 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 DeiT distilled model export workflow and publish all required artifacts to Hugging Face Hub so the fusion model and backend can load it remotely and reproducibly.

This ticket updates the workflow from:

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

Hugging Face repo: DeepFakeDetector/deit-distilled (https://huggingface.co/DeepFakeDetector/deit-distilled)

Acceptance Criteria

• A Hugging Face Write token is created and available locally as an environment variable (HF_TOKEN_WRITE)
• vit_transfer_baseline_deit.ipynb contains a final Export + Upload section
• Artifacts are exported with the standardized layout:
  • model.pth (PyTorch state_dict)
  • config.json (architecture + distilled-output handling + threshold)
  • preprocess.json (resize/crop/normalize + interpolation)
  • label_map.json (0 = real, 1 = fake)
  • README.md (model card)
• All artifacts are uploaded to the Hugging Face repo
• Notebook includes a minimal reload test (download → load → eval)
• Config explicitly documents how distilled outputs are handled (tuple logits)

Deliverables

Hugging Face (canonical storage)

Uploaded to: https://huggingface.co/DeepFakeDetector/deit-distilled

Must contain:

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

Notebook update

vit_transfer_baseline_deit.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

DeiT Distilled Output Handling

• DeiT distilled models may return tuple outputs (e.g., (logits, logits_distill))
• Backend + fusion should consume a single prob_fake, so we must define an unambiguous rule:
  • Use logits only, OR
  • Use logits_distill only, OR
  • Average both, OR
  • Weighted average (recommended default: simple average)

This rule MUST be written into config.json as distill_logits_mode.

---

Implementation Notes

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_deit.ipynb:

Cell A — Helper to unify DeiT outputs (tuple-safe)

Add this helper near the end (or as a utility cell). It ensures the same output format whether the model returns a tensor or a tuple:

import torch

def get_deit_logits(outputs, distill_mode: str = "avg"):
    """
    DeiT distilled models may return:
      - Tensor logits
      - Tuple(logits, logits_distill)
      - ModelOutput with .logits and maybe .logits_distill

    distill_mode:
      - "logits": use primary logits
      - "distill": use distillation logits
      - "avg": average both (default)
    """
    # case 1: huggingface ModelOutput
    if hasattr(outputs, "logits"):
        logits = outputs.logits
        logits_distill = getattr(outputs, "logits_distill", None)
    # case 2: tuple
    elif isinstance(outputs, (tuple, list)) and len(outputs) >= 2:
        logits, logits_distill = outputs[0], outputs[1]
    # case 3: raw tensor
    else:
        logits, logits_distill = outputs, None

    if logits_distill is None:
        return logits

    if distill_mode == "logits":
        return logits
    if distill_mode == "distill":
        return logits_distill
    # default: avg
    return 0.5 * (logits + logits_distill)

Cell B — Export artifacts (weights + configs)

import os, json, torch

EXPORT_DIR = "export/deit-distilled"
os.makedirs(EXPORT_DIR, exist_ok=True)

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

# Preprocess 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"}

# Distill output handling (must match helper above)
config = {
    "name": "deit-distilled",
    "framework": "pytorch",
    "arch": "deit_base_distilled_patch16_224",
    "patch_size": 16,
    "num_classes": 2,
    "threshold": 0.50,
    "labels": label_map,

    # CRITICAL: how to combine tuple logits
    # allowed: "logits", "distill", "avg"
    "distill_logits_mode": "avg",

    "notes": "DeiT distilled transfer 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:
  - deit
  - vision-transformer
  - deepfake-detection
  - image-classification
---

# DeiT Distilled – DeepFakeDetector

DeiT distilled (base) fine-tuned for binary deepfake detection.

## Labels
- 0 = real
- 1 = fake

## Distilled Output Handling
Some DeiT distilled implementations return two logits tensors:
- primary logits
- distillation logits

This repo stores `distill_logits_mode` in config.json (default: avg).
"""

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

print("Exported DeiT distilled artifacts to:", EXPORT_DIR)

Cell C — Upload artifacts to Hugging Face Hub

import os
from huggingface_hub import HfApi

REPO_ID = "DeepFakeDetector/deit-distilled"
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 DeiT distilled artifacts (weights + configs + preprocess)"
)

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

Cell D — (Optional but recommended) Sanity reload test (download + load)

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("DeiT distilled 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
• config.json documents distill_logits_mode used to unify tuple outputs
• Model is ready for fusion + backend consumption