Upload 2 files

audio_detector_inference.py

@@ -0,0 +1,145 @@
+"""
+audio_detector_inference.py
+============================
+Inference wrapper for AASISTDeepFake.
+Mirrors the structure of text_detector_inference.py for consistency.
+
+Usage
+-----
+from audio_detector_inference import AudioDetectorInference
+
+detector = AudioDetectorInference(checkpoint="best_aasist.pt", threshold=0.5)
+result   = detector.predict(waveform_np_array, sample_rate=16000)
+"""
+
+import os
+import numpy as np
+import torch
+from audio_model import AASISTDeepFake, SAMPLE_RATE, MAX_SAMPLES
+
+
+class AudioDetectorInference:
+    """
+    Thin wrapper around AASISTDeepFake for single-clip audio prediction.
+
+    Parameters
+    ----------
+    checkpoint : str
+        Path to the best_aasist.pt state-dict file produced by training.
+    threshold  : float
+        sigmoid(logit) >= threshold  →  Real.
+        Use 0.5 (default) or the optimal F1 threshold printed at the end
+        of training (the value labelled "Optimal threshold" in Cell 14).
+    device     : torch.device | None
+        Auto-detects CUDA if None.
+    """
+
+    def __init__(
+        self,
+        checkpoint: str   = "best_aasist.pt",
+        threshold:  float = 0.5,
+        device: torch.device = None,
+    ):
+        self.threshold = threshold
+        self.device    = device or torch.device(
+            "cuda" if torch.cuda.is_available() else "cpu"
+        )
+        self.model = None
+
+        if os.path.exists(checkpoint):
+            print(f"[AudioDetector] Loading checkpoint: {checkpoint}")
+            self.model = AASISTDeepFake()
+            self.model.load_state_dict(
+                torch.load(checkpoint, map_location=self.device)
+            )
+            self.model.eval().to(self.device)
+            print(f"[AudioDetector] ✅ AASISTDeepFake ready  "
+                  f"(threshold={self.threshold})")
+        else:
+            print(
+                f"[AudioDetector] ⚠️  '{checkpoint}' not found.\n"
+                f"[AudioDetector]    Upload best_aasist.pt to the Space — "
+                f"audio predictions will fail until then."
+            )
+
+    # ──────────────────────────────────────────────────────────────────────────
+    def _preprocess(self, x: np.ndarray, sr: int) -> torch.Tensor:
+        """
+        Normalise, convert to mono, resample to 16 kHz, and
+        pad/trim to exactly MAX_SAMPLES (80 000).
+
+        Returns
+        -------
+        torch.Tensor of shape (1, 80 000), dtype float32, on CPU.
+        """
+        x = x.astype(np.float32)
+
+        # ── Normalise int16 recordings ────────────────────────────────────────
+        if np.abs(x).max() > 1.0:
+            x = x / 32768.0
+
+        # ── Stereo → mono ─────────────────────────────────────────────────────
+        if x.ndim == 2:
+            x = x.mean(axis=1)
+
+        # ── Resample if needed ────────────────────────────────────────────────
+        if sr != SAMPLE_RATE:
+            import librosa
+            print(f"[AudioDetector] Resampling {sr} Hz → {SAMPLE_RATE} Hz …")
+            x = librosa.resample(x, orig_sr=sr, target_sr=SAMPLE_RATE)
+
+        # ── Pad or trim to exactly MAX_SAMPLES ────────────────────────────────
+        if len(x) < MAX_SAMPLES:
+            x = np.pad(x, (0, MAX_SAMPLES - len(x)))
+        else:
+            x = x[:MAX_SAMPLES]
+
+        return torch.tensor(x, dtype=torch.float32).unsqueeze(0)  # (1, 80000)
+
+    # ──────────────────────────────────────────────────────────────────────────
+    def predict(self, x: np.ndarray, sr: int) -> dict:
+        """
+        Classify a single raw audio clip.
+
+        Parameters
+        ----------
+        x  : np.ndarray   Raw waveform (any sample rate, mono or stereo).
+        sr : int          Sample rate of x.
+
+        Returns
+        -------
+        dict with keys:
+            label      : "Real" | "Fake"
+            real_prob  : P(real)  in [0, 1]
+            fake_prob  : P(fake)  in [0, 1]
+            confidence : probability of the predicted class in [0, 1]
+        """
+        if self.model is None:
+            return {
+                "error": (
+                    "Model not loaded — upload best_aasist.pt to the Space."
+                )
+            }
+
+        wav = self._preprocess(x, sr).to(self.device)
+
+        with torch.no_grad():
+            logit     = self.model(wav)                    # (1, 1)
+            real_prob = torch.sigmoid(logit).item()
+
+        fake_prob  = 1.0 - real_prob
+        is_real    = real_prob >= self.threshold
+        label      = "Real" if is_real else "Fake"
+        confidence = real_prob if is_real else fake_prob
+
+        print(
+            f"[AudioDetector] real_prob={real_prob:.4f}  "
+            f"fake_prob={fake_prob:.4f}  →  {label}"
+        )
+
+        return {
+            "label":      label,
+            "real_prob":  round(real_prob, 4),
+            "fake_prob":  round(fake_prob, 4),
+            "confidence": round(confidence, 4),
+        }

audio_model.py

@@ -0,0 +1,243 @@
+"""
+audio_model.py
+==============
+AASISTDeepFake model definition — matches the training notebook exactly.
+Import this in both training scripts and the Gradio app (via audio_detector_inference.py).
+
+Architecture:
+    Raw waveform  →  SincConv  →  Downsample (32×)  →  Res2Block
+                  →  CNN (2 layers)  →  GraphAttn (×2)  →  AttentionPool  →  Classifier
+
+Label convention (from training dataset enumerate(["fake", "real"])):
+    label = 0  →  Fake
+    label = 1  →  Real
+    sigmoid(logit) >= threshold  →  Real
+    sigmoid(logit) <  threshold  →  Fake
+"""
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# ── Audio constants (must match training) ─────────────────────────────────────
+SAMPLE_RATE  = 16_000
+MAX_DURATION = 5.0
+MAX_SAMPLES  = int(SAMPLE_RATE * MAX_DURATION)   # 80 000 samples
+
+
+# ── Sub-modules ───────────────────────────────────────────────────────────────
+
+class SincConv(nn.Module):
+    """
+    Learnable sinc-function band-pass filter bank.
+    Only 2×out_channels parameters (one f_low, one f_high per filter).
+    Initialised from mel-scale frequency bands.
+    """
+
+    @staticmethod
+    def to_mel(hz):  return 2595 * np.log10(1 + hz / 700)
+    @staticmethod
+    def to_hz(mel):  return 700 * (10 ** (mel / 2595) - 1)
+
+    def __init__(self, out_channels: int = 64, kernel_size: int = 512,
+                 sample_rate: int = 16_000):
+        super().__init__()
+        self.out_channels = out_channels
+        self.kernel_size  = kernel_size + 1 if kernel_size % 2 == 0 else kernel_size
+        self.sample_rate  = sample_rate
+
+        low_hz, high_hz = 30, sample_rate / 2 - 100
+        mel = np.linspace(self.to_mel(low_hz), self.to_mel(high_hz), out_channels + 1)
+        hz  = self.to_hz(mel)
+
+        self.low_hz_  = nn.Parameter(torch.Tensor(hz[:-1]).view(-1, 1))
+        self.band_hz_ = nn.Parameter(torch.Tensor(np.diff(hz)).view(-1, 1))
+
+        half = (self.kernel_size - 1) // 2
+        n    = torch.arange(1, half + 1, dtype=torch.float32)
+        self.register_buffer('n_',      (2 * np.pi * n / sample_rate).unsqueeze(0))
+        self.register_buffer('window_', torch.hamming_window(self.kernel_size))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        low  = 50 + torch.abs(self.low_hz_)
+        high = torch.clamp(low + 50 + torch.abs(self.band_hz_),
+                           max=self.sample_rate / 2)
+        band = (high - low)[:, 0]
+
+        f1  = torch.matmul(low,  self.n_)
+        f2  = torch.matmul(high, self.n_)
+        lp1 = torch.sin(f1) / (np.pi * self.n_ / (2 * np.pi))
+        lp2 = torch.sin(f2) / (np.pi * self.n_ / (2 * np.pi))
+        bp  = (lp2 - lp1) / (2 * band[:, None])
+
+        centre  = torch.zeros(self.out_channels, 1, device=bp.device)
+        filters = torch.cat([bp.flip(1), centre, bp], dim=1)
+        filters = filters * self.window_
+
+        x = x.unsqueeze(1)
+        return F.conv1d(x, filters.unsqueeze(1), padding=self.kernel_size // 2)
+
+
+class Res2Block(nn.Module):
+    """
+    Multi-scale residual block with inter-group accumulation.
+    Splits channels into `scale` groups; each group accumulates the previous.
+    """
+
+    def __init__(self, channels: int, scale: int = 8, dilation: int = 1):
+        super().__init__()
+        assert channels % scale == 0, \
+            f"channels ({channels}) must be divisible by scale ({scale})"
+        self.scale = scale
+        width = channels // scale
+        self.convs = nn.ModuleList([
+            nn.Conv1d(width, width, 3, padding=dilation, dilation=dilation)
+            for _ in range(scale - 1)
+        ])
+        self.bns = nn.ModuleList([nn.BatchNorm1d(width) for _ in range(scale - 1)])
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        chunks = torch.chunk(x, self.scale, dim=1)
+        out = [chunks[0]]
+        y = chunks[1]
+        for i, (conv, bn) in enumerate(zip(self.convs, self.bns)):
+            if i > 0:
+                y = y + chunks[i + 1]
+            y = F.gelu(bn(conv(y)))
+            out.append(y)
+        return torch.cat(out, dim=1)
+
+
+class GraphAttn(nn.Module):
+    """
+    Memory-efficient multi-head self-attention over temporal frames.
+    Sequences longer than 64 tokens are pooled before attention and
+    upsampled back for the residual addition.
+    """
+
+    def __init__(self, dim: int, heads: int = 4):
+        super().__init__()
+        self.heads    = heads
+        self.head_dim = dim // heads
+        self.qkv      = nn.Linear(dim, dim * 3)
+        self.out      = nn.Linear(dim, dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, N, C = x.shape
+        if N > 64:
+            x_pool = F.adaptive_avg_pool1d(
+                x.transpose(1, 2), 64).transpose(1, 2)
+        else:
+            x_pool = x
+
+        Bp, Np, Cp = x_pool.shape
+        qkv  = (self.qkv(x_pool)
+                .reshape(Bp, Np, 3, self.heads, self.head_dim)
+                .permute(2, 0, 3, 1, 4))
+        q, k, v = qkv.unbind(0)
+        attn = torch.softmax(
+            q @ k.transpose(-2, -1) / (self.head_dim ** 0.5), dim=-1)
+        out  = (attn @ v).transpose(1, 2).reshape(Bp, Np, Cp)
+        out  = self.out(out)
+
+        # Upsample back to original length for the residual connection
+        out = F.interpolate(
+            out.transpose(1, 2), size=N,
+            mode='linear', align_corners=False).transpose(1, 2)
+        return out
+
+
+class AttentionPool(nn.Module):
+    """Soft-attention weighted pooling over a sequence."""
+
+    def __init__(self, dim: int):
+        super().__init__()
+        self.attn = nn.Linear(dim, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        w = torch.softmax(self.attn(x), dim=1)   # (B, T, 1)
+        return (w * x).sum(dim=1)                # (B, dim)
+
+
+# ── Main model ────────────────────────────────────────────────────────────────
+
+class AASISTDeepFake(nn.Module):
+    """
+    AASISTDeepFake — memory-efficient raw-waveform audio spoof detector.
+
+    Input  : (B, 80 000)  float32 waveform, normalised to [-1, 1]
+    Output : (B, 1)       raw logit  →  sigmoid → P(real)
+
+    Prediction:
+        sigmoid(logit) >= threshold  →  Real  (label 1)
+        sigmoid(logit) <  threshold  →  Fake  (label 0)
+    """
+
+    def __init__(
+        self,
+        sinc_ch:     int = 64,
+        sinc_kernel: int = 512,
+        hidden:      int = 128,
+        graph_heads: int = 4,
+        n_graph:     int = 2,
+    ):
+        super().__init__()
+        self.sinc    = SincConv(sinc_ch, sinc_kernel, SAMPLE_RATE)
+        self.bn_sinc = nn.BatchNorm1d(sinc_ch)
+
+        # Aggressive downsampling: T → T/32  (kills OOM on long sequences)
+        self.downsample = nn.Sequential(
+            nn.Conv1d(sinc_ch, sinc_ch, kernel_size=8, stride=8),
+            nn.BatchNorm1d(sinc_ch), nn.GELU(),
+            nn.Conv1d(sinc_ch, sinc_ch, kernel_size=4, stride=4),
+            nn.BatchNorm1d(sinc_ch), nn.GELU(),
+        )
+        self.encoder = nn.Sequential(
+            Res2Block(sinc_ch), nn.BatchNorm1d(sinc_ch), nn.GELU(),
+        )
+        self.cnn = nn.Sequential(
+            nn.Conv1d(sinc_ch, hidden, kernel_size=3, padding=1),
+            nn.BatchNorm1d(hidden), nn.GELU(),
+            nn.Conv1d(hidden,  hidden, kernel_size=3, padding=1),
+            nn.BatchNorm1d(hidden), nn.GELU(),
+        )
+        self.graph_layers = nn.ModuleList(
+            [GraphAttn(hidden, graph_heads) for _ in range(n_graph)])
+        self.layer_norms  = nn.ModuleList(
+            [nn.LayerNorm(hidden)           for _ in range(n_graph)])
+        self.pool         = AttentionPool(hidden)
+        self.classifier   = nn.Sequential(
+            nn.LayerNorm(hidden),
+            nn.Linear(hidden, 64),
+            nn.GELU(),
+            nn.Dropout(0.4),
+            nn.Linear(64, 1),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = torch.abs(self.sinc(x))        # (B, sinc_ch, T)
+        x = F.gelu(self.bn_sinc(x))
+        x = self.downsample(x)             # (B, sinc_ch, T/32)
+        x = self.encoder(x)
+        x = self.cnn(x)                    # (B, hidden,  T/32)
+        x = x.transpose(1, 2)             # (B, T/32, hidden)
+        for attn, ln in zip(self.graph_layers, self.layer_norms):
+            x = ln(x + attn(x))
+        pooled = self.pool(x)              # (B, hidden)
+        return self.classifier(pooled)     # (B, 1)
+
+
+# ── Helper ────────────────────────────────────────────────────────────────────
+
+def load_audio_model(
+    checkpoint: str,
+    device: torch.device = None,
+) -> AASISTDeepFake:
+    """Load a trained AASISTDeepFake from a .pt state-dict checkpoint."""
+    if device is None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = AASISTDeepFake()
+    model.load_state_dict(torch.load(checkpoint, map_location=device))
+    model.eval().to(device)
+    return model