Add deepfake detection model and app files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*pth filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,19 @@
+# Python virtual environment
+venv/
+.env/
+__pycache__/
+*.pyc
+
+# OS / editor
+.DS_Store
+Thumbs.db
+.vscode/
+
+# Large Data & Models
+data/
+outputs/
+*.pth
+
+# Environment & Cache
+.venv/
+.ipynb_checkpoints/

app.py

@@ -0,0 +1,133 @@
+import gradio as gr
+
+# ---- IMPORT BACKENDS ----
+from image_backend import predict_image_pil
+# =========================
+
+
+# =========================
+# IMAGE LOGIC
+# =========================
+def analyze_image(image):
+    if image is None:
+        return "", "", "", None
+
+    label, confidence, heatmap = predict_image_pil(image)
+
+    if label == "Fake":
+        if confidence >= 90:
+            risk = '<span class="material-icons">error</span> High likelihood of deepfake'
+        elif confidence >= 60:
+            risk = '<span class="material-icons">warning</span> Possibly deepfake'
+        else:
+            risk = '<span class="material-icons">help_outline</span> Uncertain deepfake'
+    else:
+        if confidence >= 90:
+            risk = '<span class="material-icons">check_circle</span> Likely real'
+        elif confidence >= 60:
+            risk = '<span class="material-icons">warning</span> Possibly real'
+        else:
+            risk = '<span class="material-icons">help_outline</span> Uncertain – needs review'
+
+    return label, f"{confidence} %", risk, heatmap
+
+
+# =========================
+# UI
+# =========================
+with gr.Blocks() as demo:
+
+    # Load Material Icons
+    gr.Markdown("""
+    <link href="https://fonts.googleapis.com/icon?family=Material+Icons" rel="stylesheet">
+    """)
+
+    gr.Markdown("# AI Driven Deepfake Detection System")
+
+    with gr.Tabs():
+
+        # =========================
+        # HOME TAB (RESTORED)
+        # =========================
+        with gr.Tab("Home"):
+            gr.Markdown("""
+            ## Welcome
+
+            This system detects AI-generated (deepfake) content in images using
+            transformer-based deep learning models.
+            """)
+
+            gr.Markdown("""
+            ### Supported inputs
+            - Images: JPG, PNG (face-centric images recommended)
+            """)
+
+            gr.Markdown("""
+            ### How to use
+            1. Select a detection mode using the tabs above.
+            2. Upload an image or audio file.
+            3. Click **Submit** to start analysis.
+            4. Review the prediction, confidence score, and risk assessment.
+            """)
+
+            gr.Markdown("""
+            ### Understanding the results
+            - **Prediction**: Model decision (Real / Fake)
+            - **Confidence**: Certainty percentage of the prediction
+            - **Risk Assessment**:
+              - High likelihood → strong indication
+              - Possibly → caution advised
+              - Uncertain → manual review recommended
+            """)
+
+            gr.Markdown("""
+            ### Explainability
+            For images, attention heatmaps highlight the facial regions that influenced
+            the model’s decision, supporting transparency and forensic analysis.
+            """)
+
+            gr.Markdown("""
+            ### Data privacy & intended use
+            Uploaded files are processed temporarily and are not stored.
+            This system is intended as a decision-support tool and should not be used
+            as the sole source of verification.
+            """)
+
+        # =========================
+        # IMAGE TAB
+        # =========================
+        with gr.Tab("Image Deepfake"):
+            gr.Markdown("## Deepfake Image Detection")
+
+            with gr.Row():
+                with gr.Column(scale=1):
+                    image_input = gr.Image(
+                        label="Upload Image",
+                        type="pil",
+                        height=280
+                    )
+                    img_submit = gr.Button("Submit")
+                    img_clear = gr.Button("Clear")
+
+                with gr.Column(scale=2):
+                    img_pred = gr.Text(label="Prediction")
+                    img_conf = gr.Text(label="Confidence")
+                    img_risk = gr.HTML(label="Risk Assessment", value="")
+                    img_heatmap = gr.Image(
+                        label="Explainability Heatmap",
+                        height=280
+                    )
+
+            img_submit.click(
+                analyze_image,
+                image_input,
+                [img_pred, img_conf, img_risk, img_heatmap]
+            )
+
+            img_clear.click(
+                lambda: (None, "", "", "", None),
+                None,
+                [image_input, img_pred, img_conf, img_risk, img_heatmap]
+            )
+
+demo.launch(css="style.css")

evaluate.py

@@ -0,0 +1,322 @@
+import os
+import time
+import torch
+import numpy as np
+import torch.nn.functional as F
+import matplotlib.pyplot as plt
+
+from torch.utils.data import DataLoader
+from torchvision import datasets, transforms
+from transformers import ViTForImageClassification, ViTConfig
+from sklearn.metrics import (
+    accuracy_score,
+    precision_score,
+    recall_score,
+    f1_score,
+    roc_auc_score,
+    confusion_matrix,
+    roc_curve
+)
+from openpyxl import Workbook
+
+
+def main():
+    # ----------------------------------
+    # PERFORMANCE TUNING
+    # ----------------------------------
+    torch.backends.cudnn.benchmark = True
+    torch.set_num_threads(8)
+
+    # ----------------------------------
+    # Device (SAFE)
+    # ----------------------------------
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    use_amp = device.type == "cuda"
+
+    if use_amp:
+        print(f"🚀 Using device: cuda ({torch.cuda.get_device_name(0)})")
+    else:
+        print("🚀 Using device: cpu")
+
+    # ----------------------------------
+    # Paths
+    # ----------------------------------
+    base_data_dir = "data"
+
+    # 🔴 UPDATED MODEL PATH (NEW MODEL)
+    model_path = "model/vit_face_final_best.pth"
+
+    datasets_config = {
+        "FF++": f"{base_data_dir}/ff++/test",
+        "Celeb-DF": f"{base_data_dir}/celeb-df/test",
+        "DFDC": f"{base_data_dir}/dfdc/test",
+    }
+
+    os.makedirs("outputs", exist_ok=True)
+
+    # ----------------------------------
+    # Transforms
+    # ----------------------------------
+    test_tfms = transforms.Compose([
+        transforms.Resize(256),
+        transforms.CenterCrop(224),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406],
+            std=[0.229, 0.224, 0.225]
+        )
+    ])
+
+    # ----------------------------------
+    # Model
+    # ----------------------------------
+    config = ViTConfig.from_pretrained(
+        "google/vit-base-patch16-224",
+        num_labels=2
+    )
+
+    model = ViTForImageClassification.from_pretrained(
+        "google/vit-base-patch16-224",
+        config=config,
+        ignore_mismatched_sizes=True
+    )
+
+    model.load_state_dict(torch.load(model_path, map_location=device))
+    model.to(device)
+    model.eval()
+
+    print("✅ Model loaded successfully")
+
+    # ----------------------------------
+    # Evaluation
+    # ----------------------------------
+    all_results = {}
+    print("\n🧪 Running Cross-Dataset Evaluation...\n")
+
+    for ds_name, test_dir in datasets_config.items():
+        if not os.path.exists(test_dir):
+            print(f"⚠️ {ds_name}: path not found — skipping")
+            continue
+
+        test_ds = datasets.ImageFolder(test_dir, transform=test_tfms)
+
+        test_dl = DataLoader(
+            test_ds,
+            batch_size=32,
+            shuffle=False,
+            num_workers=2,      # SAFE (inside main)
+            pin_memory=True
+        )
+
+        class_to_idx = test_ds.class_to_idx
+        real_idx = class_to_idx["real"]
+        fake_idx = class_to_idx["fake"]
+
+        print(f"📂 {ds_name}: {len(test_ds)} samples | {class_to_idx}")
+
+        y_true, y_pred, y_probs = [], [], []
+        total_images = 0
+        total_time = 0.0
+
+        with torch.no_grad():
+            for imgs, labels in test_dl:
+                imgs = imgs.to(device, non_blocking=True)
+                labels = labels.to(device, non_blocking=True)
+
+                start = time.time()
+                with torch.amp.autocast(device_type="cuda", enabled=use_amp):
+                    logits = model(imgs).logits
+                end = time.time()
+
+                probs = F.softmax(logits, dim=1)[:, fake_idx]
+                preds = torch.argmax(logits, dim=1)
+
+                total_time += (end - start)
+                total_images += imgs.size(0)
+
+                y_true.extend(labels.cpu().numpy())
+                y_pred.extend(preds.cpu().numpy())
+                y_probs.extend(probs.cpu().numpy())
+
+        # ---------------- Metrics ----------------
+        acc = accuracy_score(y_true, y_pred)
+        prec = precision_score(y_true, y_pred, zero_division=0)
+        rec = recall_score(y_true, y_pred, zero_division=0)
+        f1 = f1_score(y_true, y_pred, zero_division=0)
+
+        try:
+            auc = roc_auc_score(
+                (np.array(y_true) == fake_idx).astype(int),
+                y_probs
+                )
+        except ValueError:
+            auc = float("nan")
+
+        avg_time_ms = (total_time / total_images) * 1000
+        fps = total_images / total_time
+
+        all_results[ds_name] = {
+            "acc": acc,
+            "prec": prec,
+            "rec": rec,
+            "f1": f1,
+            "auc": auc,
+            "time": avg_time_ms,
+            "fps": fps,
+        }
+
+        print(
+            f"🎯 {ds_name} | Acc: {acc:.4f} | F1: {f1:.4f} | "
+            f"AUC: {auc:.4f} | {avg_time_ms:.2f} ms/img | {fps:.2f} FPS"
+        )
+
+        # ---------------- Confusion Matrix ----------------
+        cm = confusion_matrix(y_true, y_pred, labels=[real_idx, fake_idx])
+
+        plt.figure(figsize=(5, 4))
+        plt.imshow(cm)
+        plt.title(f"{ds_name} - Confusion Matrix")
+        plt.xlabel("Predicted")
+        plt.ylabel("True")
+        plt.xticks([0, 1], ["Real", "Fake"])
+        plt.yticks([0, 1], ["Real", "Fake"])
+
+        for i in range(2):
+            for j in range(2):
+                plt.text(j, i, cm[i, j], ha="center", va="center")
+
+        plt.tight_layout()
+        plt.savefig(f"outputs/cm_{ds_name}.png")
+        plt.close()
+
+        # ---------------- ROC Curve ----------------
+        fpr, tpr, _ = roc_curve(y_true, y_probs, pos_label=fake_idx)
+
+        plt.figure(figsize=(5, 4))
+        plt.plot(fpr, tpr, label=f"AUC = {auc:.4f}")
+        plt.plot([0, 1], [0, 1], linestyle="--")
+        plt.xlabel("False Positive Rate")
+        plt.ylabel("True Positive Rate")
+        plt.title(f"{ds_name} - ROC Curve")
+        plt.legend(loc="lower right")
+        plt.tight_layout()
+        plt.savefig(f"outputs/roc_{ds_name}.png")
+        plt.close()
+
+    # ----------------------------------
+    # EXPORT (EXCEL + TEXT)
+    # ----------------------------------
+    avg_metrics = {
+        "acc": np.mean([m["acc"] for m in all_results.values()]),
+        "prec": np.mean([m["prec"] for m in all_results.values()]),
+        "rec": np.mean([m["rec"] for m in all_results.values()]),
+        "f1": np.mean([m["f1"] for m in all_results.values()]),
+        "auc": np.nanmean([m["auc"] for m in all_results.values()]),
+        "time": np.mean([m["time"] for m in all_results.values()]),
+        "fps": np.mean([m["fps"] for m in all_results.values()]),
+    }
+
+    all_results["AVERAGE"] = avg_metrics
+
+    wb = Workbook()
+    ws = wb.active
+    ws.title = "Evaluation Results"
+    ws.append(["Dataset", "Accuracy", "Precision", "Recall", "F1", "AUC", "ms/img", "FPS"])
+
+    for ds, m in all_results.items():
+        ws.append([
+            ds,
+            round(m["acc"], 4),
+            round(m["prec"], 4),
+            round(m["rec"], 4),
+            round(m["f1"], 4),
+            round(m["auc"], 4),
+            round(m["time"], 2),
+            round(m["fps"], 2),
+        ])
+
+    wb.save("outputs/evaluation_results.xlsx")
+
+    with open("outputs/summary.txt", "w") as f:
+        for ds, m in all_results.items():
+            f.write(f"Dataset: {ds}\n")
+            for k, v in m.items():
+                f.write(f"  {k.upper():<6}: {v}\n")
+            f.write("-" * 45 + "\n")
+
+    print("✅ Evaluation complete (new model evaluated successfully)")
+
+    # ----------------------------------
+    # COMBINED ROC CURVE (IEEE-FRIENDLY)
+    # ----------------------------------
+    plt.figure(figsize=(5, 4))
+
+    for ds_name in datasets_config.keys():
+        if ds_name not in all_results:
+            continue
+
+        # Reload dataset to recompute ROC cleanly
+        test_dir = datasets_config[ds_name]
+        test_ds = datasets.ImageFolder(test_dir, transform=test_tfms)
+
+        class_to_idx = test_ds.class_to_idx
+        fake_idx = class_to_idx["fake"]
+
+        test_dl = DataLoader(
+            test_ds,
+            batch_size=32,
+            shuffle=False,
+            num_workers=2,
+            pin_memory=True
+        )
+
+        y_true, y_probs = [], []
+
+        with torch.no_grad():
+            for imgs, labels in test_dl:
+                imgs = imgs.to(device, non_blocking=True)
+                labels = labels.to(device, non_blocking=True)
+
+                with torch.amp.autocast(device_type="cuda", enabled=use_amp):
+                    logits = model(imgs).logits
+
+                probs = F.softmax(logits, dim=1)[:, fake_idx]
+
+                y_true.extend((labels == fake_idx).cpu().numpy())
+                y_probs.extend(probs.cpu().numpy())
+
+        fpr, tpr, _ = roc_curve(y_true, y_probs)
+        auc_val = roc_auc_score(y_true, y_probs)
+
+        # Line styles for IEEE (print-safe)
+        if ds_name == "FF++":
+            style = "-"
+        elif ds_name == "Celeb-DF":
+            style = "--"
+        else:  # DFDC
+            style = "-."
+
+        plt.plot(
+            fpr,
+            tpr,
+            linestyle=style,
+            linewidth=2,
+            label=f"{ds_name} (AUC = {auc_val:.4f})"
+        )
+
+    # Random baseline
+    plt.plot([0, 1], [0, 1], linestyle=":", linewidth=1)
+
+    plt.xlabel("False Positive Rate")
+    plt.ylabel("True Positive Rate")
+    plt.legend(loc="lower right")
+    plt.tight_layout()
+
+    # Save IEEE-ready figure
+    plt.savefig("outputs/roc_combined.png", dpi=300, bbox_inches="tight")
+    plt.savefig("outputs/roc_combined.pdf", dpi=300, bbox_inches="tight")
+    plt.close()
+
+
+if __name__ == "__main__":
+    main()

image_backend.py

@@ -0,0 +1,89 @@
+import torch
+from torchvision import transforms
+from transformers import ViTForImageClassification, ViTConfig
+from PIL import Image
+import numpy as np
+import matplotlib.pyplot as plt
+import io
+import os
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+config = ViTConfig.from_pretrained(
+    "google/vit-base-patch16-224",
+    num_labels=2,
+    output_attentions=True
+)
+
+model = ViTForImageClassification.from_pretrained(
+    "google/vit-base-patch16-224",
+    config=config,
+    ignore_mismatched_sizes=True
+)
+
+model.load_state_dict(
+    torch.load("model/vit_face_final_best.pth", map_location=device)
+)
+
+model.to(device)
+model.eval()
+
+
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(
+        [0.485, 0.456, 0.406],
+        [0.229, 0.224, 0.225]
+    )
+])
+
+
+def get_attention_map(model, img_tensor):
+    with torch.no_grad():
+        outputs = model(img_tensor, output_attentions=True)
+        attn = outputs.attentions[-1].mean(dim=1)[0]
+        cls_attn = attn[0, 1:]
+
+        grid = int(cls_attn.size(0) ** 0.5)
+        cls_attn = cls_attn.reshape(grid, grid).cpu().numpy()
+
+        cls_attn = (cls_attn - cls_attn.min()) / (cls_attn.max() - cls_attn.min())
+        return cls_attn
+
+
+def overlay(image, heatmap):
+    heatmap = np.uint8(255 * heatmap)
+    heatmap = Image.fromarray(heatmap).resize(image.size)
+
+    fig, ax = plt.subplots(figsize=(4, 4))
+    ax.imshow(image)
+    ax.imshow(heatmap, cmap="jet", alpha=0.5)
+    ax.axis("off")
+
+    buf = io.BytesIO()
+    plt.savefig(buf, format="png", bbox_inches="tight", pad_inches=0)
+    plt.close(fig)
+    buf.seek(0)
+
+    return Image.open(buf)
+
+
+def predict_image_pil(image):
+    image = image.convert("RGB")
+
+    x = transform(image).unsqueeze(0).to(device)
+
+    with torch.no_grad():
+        outputs = model(x)
+        logits = outputs.logits
+        pred = torch.argmax(logits, dim=1).item()
+
+    label = "Fake" if pred == 0 else "Real"
+
+    heat = get_attention_map(model, x)
+    heatmap_img = overlay(image, heat)
+
+    confidence = torch.softmax(logits, dim=1)[0][pred].item() * 100
+
+    return label, round(confidence, 2), heatmap_img

model/vit_face_final_best.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7ac6c4c0cc5b00ad79871e9e44d33472a4709b0ab5fe4e3990c9882434059a5c
+size 343283314

model/vit_real_fake_best.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:296feef7a146ba07f374387161c5d04a7f043c2b7caef88fa71b97281a488f7d
+size 343283529

requirements.txt

@@ -0,0 +1,13 @@
+gradio
+torch
+torchvision
+transformers
+tensorflow
+librosa
+opencv-python
+matplotlib
+pillow
+numpy
+scikit-learn
+openpyxl
+tqdm

style.css

@@ -0,0 +1,6 @@
+.material-icons {
+    font-size: 22px;
+    color: #2c3e50;        /* Professional dark tone */
+    vertical-align: middle;
+    margin-right: 6px;
+}

train_vit_retinaface_earlystop.py

@@ -0,0 +1,270 @@
+import os
+import random
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torch.optim as optim
+
+from torch.utils.data import DataLoader, ConcatDataset
+from torchvision import datasets, transforms
+from transformers import ViTForImageClassification, ViTConfig
+from sklearn.utils.class_weight import compute_class_weight
+from sklearn.metrics import confusion_matrix, classification_report
+from tqdm import tqdm
+from multiprocessing import freeze_support
+
+# =======================================================
+# Early Stopping
+# =======================================================
+class EarlyStopping:
+    def __init__(self, patience=3, min_delta=0.001):
+        self.patience = patience
+        self.min_delta = min_delta
+        self.best_loss = None
+        self.counter = 0
+        self.stop = False
+
+    def __call__(self, val_loss):
+        if self.best_loss is None:
+            self.best_loss = val_loss
+            return
+        if val_loss > self.best_loss - self.min_delta:
+            self.counter += 1
+            print(f"⏸ EarlyStopping {self.counter}/{self.patience}")
+            if self.counter >= self.patience:
+                self.stop = True
+        else:
+            self.best_loss = val_loss
+            self.counter = 0
+
+
+def main():
+    # =======================================================
+    # Reproducibility
+    # =======================================================
+    SEED = 42
+    random.seed(SEED)
+    np.random.seed(SEED)
+    torch.manual_seed(SEED)
+    torch.cuda.manual_seed_all(SEED)
+
+    # =======================================================
+    # Device
+    # =======================================================
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    use_amp = device.type == "cuda"
+    print(f"🚀 Using device: {device}")
+
+    # =======================================================
+    # Dataset roots (YOUR STRUCTURE)
+    # =======================================================
+    DATA_DIR = "data"
+    DATASETS = {
+        "ff++": f"{DATA_DIR}/ff++",
+        "celeb-df": f"{DATA_DIR}/celeb-df",
+        "dfdc": f"{DATA_DIR}/dfdc",
+    }
+
+    # =======================================================
+    # STRONGER TRAIN AUGMENTATION (IMPORTANT)
+    # =======================================================
+    train_tfms = transforms.Compose([
+        transforms.Resize(256),
+        transforms.RandomResizedCrop(224, scale=(0.6, 1.0)),
+        transforms.RandomHorizontalFlip(),
+        transforms.RandomApply(
+            [transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)],
+            p=0.8
+        ),
+        transforms.RandomGrayscale(p=0.2),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406],
+            std=[0.229, 0.224, 0.225]
+        )
+    ])
+
+    val_tfms = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406],
+            std=[0.229, 0.224, 0.225]
+        )
+    ])
+
+    # =======================================================
+    # Load & CONCAT datasets
+    # =======================================================
+    train_sets, val_sets, test_sets = [], [], []
+
+    for name, base in DATASETS.items():
+        print(f"📂 Loading {name}")
+        train_sets.append(datasets.ImageFolder(f"{base}/train", transform=train_tfms))
+        val_sets.append(datasets.ImageFolder(f"{base}/val", transform=val_tfms))
+        test_sets.append(datasets.ImageFolder(f"{base}/test", transform=val_tfms))
+
+    train_ds = ConcatDataset(train_sets)
+    val_ds   = ConcatDataset(val_sets)
+    test_ds  = ConcatDataset(test_sets)
+
+    # =======================================================
+    # DataLoaders (IMPROVED)
+    # =======================================================
+    train_dl = DataLoader(
+        train_ds, batch_size=32, shuffle=True,
+        num_workers=0, pin_memory=True
+    )
+    val_dl = DataLoader(
+        val_ds, batch_size=32, shuffle=False,
+        num_workers=0, pin_memory=True
+    )
+    test_dl = DataLoader(
+        test_ds, batch_size=32, shuffle=False,
+        num_workers=0, pin_memory=True
+    )
+
+    classes = ["fake", "real"]
+    print("Classes:", classes)
+    print(f"Train: {len(train_ds)} | Val: {len(val_ds)} | Test: {len(test_ds)}")
+
+    # =======================================================
+    # Class Weights (BALANCED ACROSS DATASETS)
+    # =======================================================
+    all_labels = []
+    for ds in train_sets:
+        all_labels.extend([label for _, label in ds.samples])
+
+    labels_np = np.array(all_labels)
+
+    class_weights = compute_class_weight(
+        class_weight="balanced",
+        classes=np.unique(labels_np),
+        y=labels_np
+    )
+
+    class_weights = torch.tensor(class_weights, dtype=torch.float).to(device)
+    print("Class weights:", class_weights)
+
+    # =======================================================
+    # Model
+    # =======================================================
+    config = ViTConfig.from_pretrained(
+        "google/vit-base-patch16-224",
+        num_labels=2
+    )
+
+    model = ViTForImageClassification.from_pretrained(
+        "google/vit-base-patch16-224",
+        config=config,
+        ignore_mismatched_sizes=True
+    ).to(device)
+
+    # Freeze backbone
+    for param in model.vit.parameters():
+        param.requires_grad = False
+
+    # Unfreeze LAST 2 transformer blocks
+    for block in model.vit.encoder.layer[-2:]:
+        for param in block.parameters():
+            param.requires_grad = True
+
+    print("✅ Last 2 transformer blocks unfrozen")
+
+    # =======================================================
+    # Optimizer & Scheduler
+    # =======================================================
+    EPOCHS = 20
+    optimizer = optim.AdamW(model.parameters(), lr=3e-5, weight_decay=1e-4)
+    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS)
+    scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
+    early_stopper = EarlyStopping(patience=3)
+
+    # =======================================================
+    # Training Loop
+    # =======================================================
+    best_val_loss = float("inf")
+
+    for epoch in range(EPOCHS):
+        print(f"\nEpoch {epoch+1}/{EPOCHS}")
+        model.train()
+
+        correct, total = 0, 0
+
+        for imgs, labels in tqdm(train_dl):
+            imgs, labels = imgs.to(device), labels.to(device)
+            optimizer.zero_grad()
+
+            with torch.amp.autocast(device_type="cuda", enabled=use_amp):
+                logits = model(imgs).logits
+                loss = F.cross_entropy(logits, labels, weight=class_weights)
+
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
+
+            correct += (logits.argmax(1) == labels).sum().item()
+            total += labels.size(0)
+
+        train_acc = correct / total
+        print(f"Train Acc: {train_acc:.4f}")
+
+        # ---------------- Validation ----------------
+        model.eval()
+        val_loss, val_correct, val_total = 0.0, 0, 0
+
+        with torch.no_grad():
+            for imgs, labels in val_dl:
+                imgs, labels = imgs.to(device), labels.to(device)
+                with torch.amp.autocast(device_type="cuda", enabled=use_amp):
+                    logits = model(imgs).logits
+                    loss = F.cross_entropy(logits, labels, weight=class_weights)
+
+                val_loss += loss.item()
+                val_correct += (logits.argmax(1) == labels).sum().item()
+                val_total += labels.size(0)
+
+        val_loss /= len(val_dl)
+        val_acc = val_correct / val_total
+        print(f"Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.4f}")
+
+        scheduler.step()
+        early_stopper(val_loss)
+
+        if val_loss < best_val_loss:
+            best_val_loss = val_loss
+            torch.save(model.state_dict(), "vit_face_final_best.pth")
+            print("✅ Best model saved")
+
+        if early_stopper.stop:
+            print("🛑 Early stopping triggered")
+            break
+
+    # =======================================================
+    # Testing (AMP ENABLED)
+    # =======================================================
+    model.eval()
+    all_preds, all_labels = [], []
+
+    with torch.no_grad():
+        for imgs, labels in test_dl:
+            imgs = imgs.to(device)
+            with torch.amp.autocast(device_type="cuda", enabled=use_amp):
+                logits = model(imgs).logits
+
+            preds = logits.argmax(1).cpu().numpy()
+            all_preds.extend(preds)
+            all_labels.extend(labels.numpy())
+
+    print("\n📊 Confusion Matrix")
+    print(confusion_matrix(all_labels, all_preds))
+
+    print("\n📊 Classification Report")
+    print(classification_report(all_labels, all_preds, target_names=classes))
+
+    print("✅ FINAL TRAINING COMPLETED SUCCESSFULLY")
+
+
+if __name__ == "__main__":
+    freeze_support()
+    main()