soHARSH/AASIST

1# =========================2# IMPORTS3# =========================4import gradio as gr5import torch6import torch.nn as nn7import torch.nn.functional as F8import numpy as np9import librosa10 11# =========================12# ATTENTION (REQUIRED ✅)13# =========================14class SelfAttention(nn.Module):15    def __init__(self, in_dim):16        super().__init__()17        self.query = nn.Conv1d(in_dim, in_dim // 8, 1)18        self.key = nn.Conv1d(in_dim, in_dim // 8, 1)19        self.value = nn.Conv1d(in_dim, in_dim, 1)20        self.gamma = nn.Parameter(torch.zeros(1))21 22    def forward(self, x):23        Q = self.query(x)24        K = self.key(x)25        V = self.value(x)26 27        attention = torch.bmm(Q.permute(0, 2, 1), K)28        attention = F.softmax(attention, dim=-1)29 30        out = torch.bmm(V, attention.permute(0, 2, 1))31        return self.gamma * out + x32 33 34# =========================35# MODEL (MATCH TRAINED)36# =========================37class SpoofModel(nn.Module):38    def __init__(self):39        super().__init__()40 41        # ✅ SAME NAME: cnn42        self.cnn = nn.Sequential(43            nn.Conv2d(1, 32, 3, padding=1),44            nn.BatchNorm2d(32),45            nn.ReLU(),46            nn.MaxPool2d(2),47 48            nn.Conv2d(32, 64, 3, padding=1),49            nn.BatchNorm2d(64),50            nn.ReLU(),51            nn.MaxPool2d(2),52 53            nn.Conv2d(64, 128, 3, padding=1),54            nn.BatchNorm2d(128),55            nn.ReLU(),56            nn.MaxPool2d(2),57        )58 59        # ✅ ATTENTION PRESENT60        self.attention = SelfAttention(128)61 62        # ✅ FC SIZE = 6400 (FROM ERROR)63        self.fc = nn.Sequential(64            nn.Linear(6400, 256),65            nn.ReLU(),66            nn.Dropout(0.5),67            nn.Linear(256, 2)68        )69 70    def forward(self, x):71        x = self.cnn(x)72 73        # collapse frequency74        x = torch.mean(x, dim=2)75 76        # attention77        x = self.attention(x)78 79        x = x.view(x.size(0), -1)80        return self.fc(x)81 82 83# =========================84# LOAD MODEL85# =========================86device = torch.device("cpu")87 88model = SpoofModel()89 90state_dict = torch.load("model.pth", map_location=device)91model.load_state_dict(state_dict)92 93model.to(device)94model.eval()95 96print("✅ Model loaded successfully")97 98 99# =========================100# PREPROCESS (IMPORTANT ⚠️)101# =========================102def preprocess(audio):103    sr, y = audio104 105    y = y.astype(np.float32)106 107    # stereo → mono108    if len(y.shape) > 1:109        y = np.mean(y, axis=1)110 111    # ⚠️ SAME AS TRAINING112    mel = librosa.feature.melspectrogram(113        y=y,114        sr=sr,115        n_mels=64116    )117 118    feat = librosa.power_to_db(mel)119 120    feat = (feat - np.mean(feat)) / (np.std(feat) + 1e-6)121 122    # ⚠️ LENGTH MUST MATCH TRAINING123    max_len = 400124    if feat.shape[1] < max_len:125        feat = np.pad(feat, ((0, 0), (0, max_len - feat.shape[1])))126    else:127        feat = feat[:, :max_len]128 129    feat = torch.tensor(feat).unsqueeze(0).unsqueeze(0).float()130 131    return feat132 133 134# =========================135# PREDICT136# =========================137def predict(audio):138    if audio is None:139        return {"Error": 1.0}140 141    X = preprocess(audio).to(device)142 143    with torch.no_grad():144        output = model(X)145        probs = F.softmax(output, dim=1)146 147    return {148        "Bonafide (Real)": float(probs[0][1]),149        "Spoof (Fake)": float(probs[0][0])150    }151 152 153# =========================154# UI155# =========================156demo = gr.Interface(157    fn=predict,158    inputs=gr.Audio(type="numpy", label="🎤 Upload or Record Audio"),159    outputs=gr.Label(label="🧠 Prediction"),160    title="🎙️ Voice Deepfake Detection",161    description="Detect whether audio is REAL or FAKE",162    theme="soft"163)164 165demo.launch()