sync from GitHub (0154d02)

audio/aasist3/.gitignore

@@ -0,0 +1 @@
+weights/

audio/aasist3/Dockerfile

@@ -0,0 +1,35 @@
+FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
+
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PYTHONUNBUFFERED=1
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    python3 python3-pip \
+    ffmpeg libsndfile1 \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN ln -sf /usr/bin/python3 /usr/bin/python
+
+WORKDIR /app
+
+# Install PyTorch with CUDA 12.1
+RUN pip install --no-cache-dir \
+    torch==2.5.1 torchaudio==2.5.1 \
+    --index-url https://download.pytorch.org/whl/cu121
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+RUN python -c "from transformers import Wav2Vec2Config; Wav2Vec2Config.from_pretrained('facebook/wav2vec2-large-xlsr-53', cache_dir='/app/w2v_cache')"
+
+COPY model/ /app/model/
+COPY api.py .
+RUN mkdir -p /app/weights
+COPY weights/ /app/weights/
+
+EXPOSE 8005
+
+RUN adduser --disabled-password --gecos '' appuser
+USER appuser
+
+CMD ["python", "api.py"]

audio/aasist3/api.py

@@ -0,0 +1,345 @@
+"""AASIST3 Audio Deepfake Detection API.
+
+Detects synthetic speech using the AASIST3 model architecture:
+- Frontend: XLSR wav2vec 2.0 (HuggingFace Transformers)
+- Backend: AASIST with KAN (Kolmogorov-Arnold Network) linear
+  layers and Graph Attention Networks
+
+Reference: https://github.com/AI4Bharat/AASIST3
+"""
+
+import base64
+import io
+import logging
+import os
+import platform
+import time
+from typing import Optional
+
+import numpy as np
+import soundfile as sf
+import torch
+import torch.nn.functional as F
+import uvicorn
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field
+
+# Point transformers cache to pre-cached wav2vec2 config
+# (must be set before importing model code)
+os.environ["TRANSFORMERS_CACHE"] = "/app/w2v_cache"
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger("aasist3_api")
+
+# Import model class
+try:
+    from model import aasist3 as AASIST3Model
+except ImportError as e:
+    logger.error(f"Failed to import AASIST3 model: {e}")
+    AASIST3Model = None
+
+# Constants
+MODEL_NAME = "aasist3"
+MODEL_ID = "aasist3_kan_mlaad"
+WEIGHTS_DIR = "/app/weights"
+
+
+def _get_device():
+    """Select optimal device: MPS (Apple) > CUDA (NVIDIA) > CPU."""
+    override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower()
+    if override == "cpu":
+        return torch.device("cpu")
+    if override == "cuda" and torch.cuda.is_available():
+        return torch.device("cuda")
+    if (
+        override == "mps"
+        and hasattr(torch.backends, "mps")
+        and torch.backends.mps.is_available()
+    ):
+        return torch.device("mps")
+    if override:
+        pass  # Invalid override, fall through to auto-detect
+    if (
+        platform.system() == "Darwin"
+        and hasattr(torch.backends, "mps")
+        and torch.backends.mps.is_available()
+    ):
+        return torch.device("mps")
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    return torch.device("cpu")
+
+
+DEVICE = _get_device()
+
+if DEVICE.type == "cuda":
+    torch.backends.cudnn.benchmark = True
+    torch.set_float32_matmul_precision("high")
+
+if DEVICE.type == "cuda":
+    logger.info(
+        "Device: cuda (%s, %.1f GB VRAM)",
+        torch.cuda.get_device_name(0),
+        torch.cuda.get_device_properties(0).total_memory / 1024**3,
+    )
+else:
+    logger.warning(
+        "Device: %s (no CUDA available -- check nvidia-container-toolkit)",
+        DEVICE,
+    )
+
+SAMPLE_RATE = 16000
+TARGET_SAMPLES = 64600  # ~4.04 seconds at 16kHz
+
+# Global model instance
+model = None
+
+
+class AudioInput(BaseModel):
+    """Request schema for audio deepfake detection."""
+
+    audio_data: str = Field(
+        ..., description="Base64 encoded audio string (WAV/MP3/etc)"
+    )
+    threshold: Optional[float] = Field(
+        0.5, ge=0.0, le=1.0, description="Classification threshold"
+    )
+
+
+app = FastAPI(
+    title="AASIST3 Audio Deepfake Detection API",
+    description=(
+        "Service for detecting synthetic speech using the "
+        "AASIST3 model (HuggingFace wav2vec 2.0 + AASIST "
+        "with KAN layers)."
+    ),
+    version="1.0.0",
+)
+
+
+def load_model():
+    """Load the AASIST3 model from pretrained weights.
+
+    Returns:
+        The loaded model, or None if loading fails.
+    """
+    global model
+    if model is not None:
+        return model
+
+    logger.info(f"Loading AASIST3 model onto {DEVICE}...")
+
+    if AASIST3Model is None:
+        logger.error("AASIST3 model class not available.")
+        return None
+
+    weights_safetensors = os.path.join(WEIGHTS_DIR, "model.safetensors")
+    weights_config = os.path.join(WEIGHTS_DIR, "config.json")
+
+    if not os.path.exists(weights_safetensors):
+        logger.error(f"Model weights not found at {weights_safetensors}")
+        return None
+
+    if not os.path.exists(weights_config):
+        logger.error(f"Model config not found at {weights_config}")
+        return None
+
+    try:
+        model = AASIST3Model.from_pretrained(WEIGHTS_DIR)
+        model.to(DEVICE)
+        # Set model to inference mode (disables dropout, batchnorm)
+        model.train(False)
+
+        logger.info("AASIST3 model loaded successfully.")
+        return model
+    except Exception as e:
+        logger.exception(f"Failed to load AASIST3 model: {e}")
+        model = None
+        return None
+
+
+@app.on_event("startup")
+async def startup_event():
+    """Load model on service startup."""
+    load_model()
+
+
+def _gpu_health_info() -> dict:
+    """Return GPU metrics for the health endpoint."""
+    if torch.cuda.is_available() and DEVICE.type == "cuda":
+        return {
+            "gpu_name": torch.cuda.get_device_name(0),
+            "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2),
+            "vram_total_mb": round(
+                torch.cuda.get_device_properties(0).total_memory / 1024**2
+            ),
+        }
+    return {}
+
+
+@app.get("/health")
+async def health():
+    """Health check endpoint."""
+    return {
+        "status": "healthy" if model is not None else "degraded",
+        "model": MODEL_NAME,
+        "model_id": MODEL_ID,
+        "device": str(DEVICE),
+        "weights_found": os.path.exists(os.path.join(WEIGHTS_DIR, "model.safetensors")),
+        **_gpu_health_info(),
+    }
+
+
+def _load_audio_bytes(audio_bytes: bytes) -> tuple:
+    """Load audio from raw bytes using soundfile with torchaudio fallback.
+
+    Args:
+        audio_bytes: Raw audio file bytes.
+
+    Returns:
+        Tuple of (audio_numpy_array, sample_rate).
+
+    Raises:
+        ValueError: If audio cannot be loaded by any backend.
+    """
+    # Try soundfile first (handles WAV, FLAC natively)
+    sf_error = None
+    try:
+        audio, sr = sf.read(io.BytesIO(audio_bytes), dtype="float32")
+        if audio.ndim > 1:
+            audio = audio.mean(axis=1)  # Convert to mono
+        return audio, sr
+    except Exception as sf_err:
+        sf_error = sf_err
+        logger.debug(f"soundfile failed, trying torchaudio: {sf_err}")
+
+    # Fallback to torchaudio (handles MP3, compressed formats)
+    try:
+        import torchaudio
+
+        buf = io.BytesIO(audio_bytes)
+        waveform, sr = torchaudio.load(buf)
+        if waveform.shape[0] > 1:
+            waveform = waveform.mean(dim=0, keepdim=True)
+        return waveform.squeeze(0).numpy(), sr
+    except Exception as ta_err:
+        raise ValueError(
+            f"Failed to load audio with soundfile and torchaudio: "
+            f"sf={sf_error}, ta={ta_err}"
+        )
+
+
+def preprocess_audio(audio_bytes: bytes) -> torch.Tensor:
+    """Preprocess audio for AASIST3 inference.
+
+    Loads audio, resamples to 16kHz mono, and zero-pads or
+    truncates to TARGET_SAMPLES.
+
+    Args:
+        audio_bytes: Raw audio file bytes.
+
+    Returns:
+        Audio tensor of shape (1, TARGET_SAMPLES).
+
+    Raises:
+        ValueError: If audio preprocessing fails.
+    """
+    try:
+        logger.info("Starting audio preprocessing...")
+        audio, sr = _load_audio_bytes(audio_bytes)
+        logger.info(f"Audio loaded. Length: {len(audio)} samples at {sr}Hz")
+
+        # Resample to 16kHz if needed
+        if sr != SAMPLE_RATE:
+            import torchaudio
+
+            resampler = torchaudio.transforms.Resample(
+                orig_freq=sr, new_freq=SAMPLE_RATE
+            )
+            audio_tensor = torch.FloatTensor(audio).unsqueeze(0)
+            audio_tensor = resampler(audio_tensor).squeeze(0)
+            audio = audio_tensor.numpy()
+            logger.info(
+                f"Resampled from {sr}Hz to {SAMPLE_RATE}Hz. "
+                f"New length: {len(audio)} samples"
+            )
+
+        # Zero-pad or truncate to TARGET_SAMPLES
+        if len(audio) >= TARGET_SAMPLES:
+            audio = audio[:TARGET_SAMPLES]
+        else:
+            pad_length = TARGET_SAMPLES - len(audio)
+            audio = np.pad(audio, (0, pad_length), mode="constant")
+
+        logger.info(f"Audio padded/trimmed to {TARGET_SAMPLES} samples")
+
+        audio_tensor = torch.FloatTensor(audio).unsqueeze(0).to(DEVICE)
+        return audio_tensor
+    except Exception as e:
+        logger.error(f"Error preprocessing audio: {e}")
+        raise ValueError(f"Audio preprocessing failed: {str(e)}")
+
+
+@app.post("/predict")
+async def predict(input_data: AudioInput):
+    """Run deepfake detection on base64-encoded audio.
+
+    The model outputs 2 logits: [bonafide_score, spoof_score].
+    Class 0 = bonafide (real), Class 1 = spoof (fake).
+    The returned probability is the spoof/fake probability.
+    """
+    if model is None:
+        if load_model() is None:
+            raise HTTPException(status_code=503, detail="Model not loaded")
+
+    try:
+        start_time = time.time()
+        logger.info(
+            f"Prediction request. Data size: " f"{len(input_data.audio_data)} chars"
+        )
+
+        # Decode base64 audio
+        audio_bytes = base64.b64decode(input_data.audio_data)
+
+        # Preprocess
+        audio_tensor = preprocess_audio(audio_bytes)
+
+        # Inference
+        logger.info("Starting model inference...")
+        with torch.no_grad():
+            output = model(audio_tensor)
+
+            # output shape: [batch, 2]
+            # Index 0 = bonafide logit, Index 1 = spoof logit
+            probs = torch.softmax(output, dim=1)
+            prob_fake = probs[0, 1].item()
+
+        prediction = 1 if prob_fake >= input_data.threshold else 0
+        verdict = "fake" if prediction == 1 else "real"
+        inference_time = time.time() - start_time
+
+        logger.info(
+            f"Prediction: {verdict} (prob_fake={prob_fake:.4f}, "
+            f"time={inference_time:.3f}s)"
+        )
+
+        return {
+            "model": MODEL_NAME,
+            "probability": float(prob_fake),
+            "prediction": int(prediction),
+            "class": verdict,
+            "inference_time": float(inference_time),
+        }
+
+    except Exception as e:
+        logger.exception(f"Error during prediction: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+if __name__ == "__main__":
+    port = int(os.environ.get("MODEL_PORT", 8005))
+    uvicorn.run(app, host="0.0.0.0", port=port)

audio/aasist3/model/__init__.py

@@ -0,0 +1 @@
+from .full_model import aasist3

audio/aasist3/model/branch.py

@@ -0,0 +1,34 @@
+import torch.nn as nn
+
+from .hs_gal import HtrgGraphAttentionLayer
+from .pool import GraphPool
+
+class InferenceBranch(nn.Module):
+    def __init__(self, gat_dims, temperature, pool_ratio, size):
+        super().__init__()
+        self.htrg_gat1 = HtrgGraphAttentionLayer(
+            gat_dims[0], gat_dims[1], temperature=temperature, size=size
+        )
+        self.htrg_gat2 = HtrgGraphAttentionLayer(
+            gat_dims[1], gat_dims[1], temperature=temperature, size=size
+        )
+
+        self.pool_hS = GraphPool(pool_ratio, gat_dims[1], 0.3, size=size)
+        self.pool_hT = GraphPool(pool_ratio, gat_dims[1], 0.3, size=size)
+
+    def forward(self, out_T, out_S, master):
+        # Первая стадия
+        out_T_res, out_S_res, master_res = self.htrg_gat1(out_T, out_S, master=master)
+
+        # Пулинг
+        out_S_res = self.pool_hS(out_S_res)
+        out_T_res = self.pool_hT(out_T_res)
+
+        # Вторая стадия с residual connection
+        out_T_aug, out_S_aug, master_aug = self.htrg_gat2(out_T_res, out_S_res, master=master_res)
+        
+        out_T_final = out_T_res + out_T_aug
+        out_S_final = out_S_res + out_S_aug
+        master_final = master_res + master_aug
+
+        return out_T_final, out_S_final, master_final

audio/aasist3/model/full_model.py

@@ -0,0 +1,139 @@
+#!/usr/bin/env python
+
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+from huggingface_hub import PyTorchModelHubMixin
+
+from .kan import KANLinear
+from .gat import GraphAttentionLayer
+from .pool import GraphPool
+from .branch import InferenceBranch
+from .residual import Residual_block
+from .wav2vec import Wav2Vec2Encoder
+
+
+class aasist3(nn.Module, PyTorchModelHubMixin):
+    def __init__(self, d_args={
+        "architecture": "AASIST",
+        "nb_samp": 64600,
+        "first_conv": 128,
+        "filts": [70, [1, 32], [32, 32], [32, 64], [64, 64]],
+        "gat_dims": [64, 32],
+        "pool_ratios": [0.5, 0.7, 0.5, 0.5],
+        "temperatures": [2.0, 2.0, 100.0, 100.0],
+    }, size=200, w2v_cache_dir="weights/", load_pretrained=True):
+        super().__init__()
+
+        self.w2v_encoder = Wav2Vec2Encoder(cache_dir=w2v_cache_dir, load_pretrained=load_pretrained)
+        self.bridge = KANLinear(1024, 128)
+
+        self.d_args = d_args
+        filts = d_args["filts"]
+        gat_dims = d_args["gat_dims"]
+        pool_ratios = d_args["pool_ratios"]
+        temperatures = d_args["temperatures"]
+
+        self.first_bn = nn.BatchNorm2d(num_features=1)
+        self.selu = nn.SELU(inplace=True)
+        self.drop = nn.Dropout(0.5, inplace=True)
+        self.drop_way = nn.Dropout(0.2, inplace=True)
+
+        self.encoder = nn.Sequential(
+            nn.Sequential(Residual_block(nb_filts=filts[1], first=True)),
+            nn.Sequential(Residual_block(nb_filts=filts[2])),
+            nn.Sequential(Residual_block(nb_filts=filts[3])),
+            nn.Sequential(Residual_block(nb_filts=filts[4])),
+            nn.Sequential(Residual_block(nb_filts=filts[4])),
+            nn.Sequential(Residual_block(nb_filts=filts[4])))
+
+        self.pos_S = nn.Parameter(torch.randn(1, 42, filts[-1][-1]))
+        self.pos_T = nn.Parameter(torch.randn(1, 67, filts[-1][-1]))
+
+        self.GAT_layer_S = GraphAttentionLayer(filts[-1][-1], gat_dims[0], temperature=temperatures[0], size=size)
+        self.GAT_layer_T = GraphAttentionLayer(filts[-1][-1], gat_dims[0], temperature=temperatures[1], size=size)
+
+        self.pool_S = GraphPool(pool_ratios[0], gat_dims[0], 0.3, size=size)
+        self.pool_T = GraphPool(pool_ratios[1], gat_dims[0], 0.3, size=size)
+
+        self.master1 = nn.Parameter(torch.randn(1, 1, gat_dims[0]))
+        self.master2 = nn.Parameter(torch.randn(1, 1, gat_dims[0]))
+        self.master3 = nn.Parameter(torch.randn(1, 1, gat_dims[0]))
+        self.master4 = nn.Parameter(torch.randn(1, 1, gat_dims[0]))
+
+        self.inference_branch1 = InferenceBranch(
+            gat_dims=gat_dims,
+            temperature=temperatures[2],
+            pool_ratio=pool_ratios[2],
+            size=size
+        )
+        self.inference_branch2 = InferenceBranch(
+            gat_dims=gat_dims,
+            temperature=temperatures[2],
+            pool_ratio=pool_ratios[2],
+            size=size
+        )
+        self.inference_branch3 = InferenceBranch(
+            gat_dims=gat_dims,
+            temperature=temperatures[2],
+            pool_ratio=pool_ratios[2],
+            size=size
+        )
+        self.inference_branch4 = InferenceBranch(
+            gat_dims=gat_dims,
+            temperature=temperatures[2],
+            pool_ratio=pool_ratios[2],
+            size=size
+        )
+
+        self.out_layer = KANLinear(5 * gat_dims[1], 2)
+
+    def forward(self, x, Freq_aug=False):
+        x = self.w2v_encoder(x)
+        x = self.bridge(x)
+        x = x.transpose(1, 2)
+        x = x.unsqueeze(dim=1)
+        x = F.max_pool2d(torch.abs(x), (3, 3))
+        x = self.first_bn(x)
+        x = self.selu(x)
+
+        e = self.encoder(x)
+
+        # GAT-S
+        e_S, _ = torch.max(torch.abs(e), dim=3)
+        e_S = e_S.transpose(1, 2) + self.pos_S
+        gat_S = self.GAT_layer_S(e_S)
+        out_S = self.pool_S(gat_S)
+
+        # GAT-T
+        e_T, _ = torch.max(torch.abs(e), dim=2)
+        e_T = e_T.transpose(1, 2) + self.pos_T
+        gat_T = self.GAT_layer_T(e_T)
+        out_T = self.pool_T(gat_T)
+
+        out_T1, out_S1, master1 = self.inference_branch1(out_T, out_S, self.master1)
+        out_T2, out_S2, master2 = self.inference_branch2(out_T, out_S, self.master2)
+        out_T3, out_S3, master3 = self.inference_branch3(out_T, out_S, self.master3)
+        out_T4, out_S4, master4 = self.inference_branch4(out_T, out_S, self.master4)
+
+        out_T1, out_T2 = self.drop_way(out_T1), self.drop_way(out_T2)
+        out_T3, out_T4 = self.drop_way(out_T3), self.drop_way(out_T4)
+        out_S1, out_S2 = self.drop_way(out_S1), self.drop_way(out_S2)
+        out_S3, out_S4 = self.drop_way(out_S3), self.drop_way(out_S4)
+        master1, master2 = self.drop_way(master1), self.drop_way(master2)
+        master3, master4 = self.drop_way(master3), self.drop_way(master4)
+
+        out_T = torch.stack([out_T1, out_T2, out_T3, out_T4]).max(dim=0)[0]
+        out_S = torch.stack([out_S1, out_S2, out_S3, out_S4]).max(dim=0)[0]
+        master = torch.stack([master1, master2, master3, master4]).max(dim=0)[0]
+
+        T_max, _ = torch.max(torch.abs(out_T), dim=1)
+        T_avg = torch.mean(out_T, dim=1)
+        S_max, _ = torch.max(torch.abs(out_S), dim=1)
+        S_avg = torch.mean(out_S, dim=1)
+
+        last_hidden = torch.cat([T_max, T_avg, S_max, S_avg, master.squeeze(1)], dim=1)
+        last_hidden = self.drop(last_hidden)
+        output = self.out_layer(last_hidden)
+
+        return output

audio/aasist3/model/gat.py

@@ -0,0 +1,99 @@
+import torch, torch.nn as nn, torch.nn.functional as F
+
+from .kan import KANLinear
+
+
+class GraphAttentionLayer(nn.Module):
+    def __init__(self, in_dim, out_dim, size, **kwargs):
+        super().__init__()
+
+        # attention map
+        self.att_proj = KANLinear(in_dim, out_dim)
+        self.att_weight = self._init_new_params(out_dim, 1)
+
+        # project
+        self.proj_with_att = KANLinear(in_dim, out_dim)
+        self.proj_without_att = KANLinear(in_dim, out_dim)
+
+        # batch norm
+        self.bn = nn.BatchNorm1d(out_dim)
+
+        # dropout for inputs
+        self.input_drop = nn.Dropout(p=0.2)
+
+        # activate
+        self.act = nn.SELU(inplace=True)
+
+        # temperature
+        self.temp = 1.
+        if "temperature" in kwargs:
+            self.temp = kwargs["temperature"]
+
+    def forward(self, x):
+        '''
+        x   :(#bs, #node, #dim)
+        '''
+        # apply input dropout
+        x = self.input_drop(x)
+
+        # derive attention map
+        att_map = self._derive_att_map(x)
+
+        # projection
+        x = self._project(x, att_map)
+
+        # apply batch norm
+        x = self._apply_BN(x)
+        x = self.act(x)
+        return x
+
+    def _pairwise_mul_nodes(self, x):
+        '''
+        Calculates pairwise multiplication of nodes.
+        - for attention map
+        x           :(#bs, #node, #dim)
+        out_shape   :(#bs, #node, #node, #dim)
+        '''
+
+        nb_nodes = x.size(1)
+        x = x.unsqueeze(2).expand(-1, -1, nb_nodes, -1)
+        x_mirror = x.transpose(1, 2)
+
+        return x * x_mirror
+
+    def _derive_att_map(self, x):
+        '''
+        x           :(#bs, #node, #dim)
+        out_shape   :(#bs, #node, #node, 1)
+        '''
+        att_map = self._pairwise_mul_nodes(x)
+        # size: (#bs, #node, #node, #dim_out)
+        att_map = torch.tanh(self.att_proj(att_map))
+        # size: (#bs, #node, #node, 1)
+        att_map = torch.matmul(att_map, self.att_weight)
+
+        # apply temperature
+        att_map = att_map / self.temp
+
+        att_map = F.softmax(att_map, dim=-2)
+
+        return att_map
+
+    def _project(self, x, att_map):
+        x1 = self.proj_with_att(torch.matmul(att_map.squeeze(-1), x))
+        x2 = self.proj_without_att(x)
+
+        return x1 + x2
+
+    def _apply_BN(self, x):
+        org_size = x.size()
+        x = x.view(-1, org_size[-1])
+        x = self.bn(x)
+        x = x.view(org_size)
+
+        return x
+
+    def _init_new_params(self, *size):
+        out = nn.Parameter(torch.FloatTensor(*size))
+        nn.init.xavier_normal_(out)
+        return out

audio/aasist3/model/hs_gal.py

@@ -0,0 +1,176 @@
+import torch, torch.nn as nn, torch.nn.functional as F
+
+from .kan import KANLinear
+
+
+class HtrgGraphAttentionLayer(nn.Module):
+    def __init__(self, in_dim, out_dim, size, **kwargs):
+        super().__init__()
+
+        self.proj_type1 = KANLinear(in_dim, in_dim)
+        self.proj_type2 = KANLinear(in_dim, in_dim)
+
+        # attention map
+        self.att_proj = KANLinear(in_dim, out_dim)
+        self.att_projM = KANLinear(in_dim, out_dim)
+
+        self.att_weight11 = self._init_new_params(out_dim, 1)
+        self.att_weight22 = self._init_new_params(out_dim, 1)
+        self.att_weight12 = self._init_new_params(out_dim, 1)
+        self.att_weightM = self._init_new_params(out_dim, 1)
+
+        # project
+        self.proj_with_att = KANLinear(in_dim, out_dim)
+        self.proj_without_att = KANLinear(in_dim, out_dim)
+
+        self.proj_with_attM = KANLinear(in_dim, out_dim)
+        self.proj_without_attM = KANLinear(in_dim, out_dim)
+
+        # batch norm
+        self.bn = nn.BatchNorm1d(out_dim)
+
+        # dropout for inputs
+        self.input_drop = nn.Dropout(p=0.2)
+
+        # activate
+        self.act = nn.SELU(inplace=True)
+
+        # temperature
+        self.temp = 1.
+        if "temperature" in kwargs:
+            self.temp = kwargs["temperature"]
+
+    def forward(self, x1, x2, master=None):
+        '''
+        x1  :(#bs, #node, #dim)
+        x2  :(#bs, #node, #dim)
+        '''
+        num_type1 = x1.size(1)
+        num_type2 = x2.size(1)
+
+        x1 = self.proj_type1(x1)
+        x2 = self.proj_type2(x2)
+
+        x = torch.cat([x1, x2], dim=1)
+
+        if master is None:
+            master = torch.mean(x, dim=1, keepdim=True)
+
+        # apply input dropout
+        x = self.input_drop(x)
+
+        # derive attention map
+        att_map = self._derive_att_map(x, num_type1, num_type2)
+
+        # directional edge for master node
+        master = self._update_master(x, master)
+
+        # projection
+        x = self._project(x, att_map)
+
+        # apply batch norm
+        x = self._apply_BN(x)
+        # x = self.act(x)
+
+        x1 = x.narrow(1, 0, num_type1)
+        x2 = x.narrow(1, num_type1, num_type2)
+
+        return x1, x2, master
+
+    def _update_master(self, x, master):
+
+        att_map = self._derive_att_map_master(x, master)
+        master = self._project_master(x, master, att_map)
+
+        return master
+
+    def _pairwise_mul_nodes(self, x):
+        '''
+        Calculates pairwise multiplication of nodes.
+        - for attention map
+        x           :(#bs, #node, #dim)
+        out_shape   :(#bs, #node, #node, #dim)
+        '''
+
+        nb_nodes = x.size(1)
+        x = x.unsqueeze(2).expand(-1, -1, nb_nodes, -1)
+        x_mirror = x.transpose(1, 2)
+
+        return x * x_mirror
+
+    def _derive_att_map_master(self, x, master):
+        '''
+        x           :(#bs, #node, #dim)
+        out_shape   :(#bs, #node, #node, 1)
+        '''
+        att_map = x * master
+        att_map = torch.tanh(self.att_projM(att_map))
+
+        att_map = torch.matmul(att_map, self.att_weightM)
+
+        # apply temperature
+        att_map = att_map / self.temp
+
+        att_map = F.softmax(att_map, dim=-2)
+
+        return att_map
+
+    def _derive_att_map(self, x, num_type1, num_type2):
+        '''
+        x           :(#bs, #node, #dim)
+        out_shape   :(#bs, #node, #node, 1)
+        '''
+        att_map = self._pairwise_mul_nodes(x)
+        # size: (#bs, #node, #node, #dim_out)
+        att_map = torch.tanh(self.att_proj(att_map))
+        # size: (#bs, #node, #node, 1)
+
+        att_board = torch.zeros_like(att_map[:, :, :, 0]).unsqueeze(-1)
+
+        att_board[:, :num_type1, :num_type1, :] = torch.matmul(
+            att_map[:, :num_type1, :num_type1, :], self.att_weight11)
+        att_board[:, num_type1:, num_type1:, :] = torch.matmul(
+            att_map[:, num_type1:, num_type1:, :], self.att_weight22)
+        att_board[:, :num_type1, num_type1:, :] = torch.matmul(
+            att_map[:, :num_type1, num_type1:, :], self.att_weight12)
+        att_board[:, num_type1:, :num_type1, :] = torch.matmul(
+            att_map[:, num_type1:, :num_type1, :], self.att_weight12)
+
+        att_map = att_board
+
+        # att_map = torch.matmul(att_map, self.att_weight12)
+
+        # apply temperature
+        att_map = att_map / self.temp
+
+        att_map = F.softmax(att_map, dim=-2)
+
+        return att_map
+
+    def _project(self, x, att_map):
+        x1 = self.proj_with_att(torch.matmul(att_map.squeeze(-1), x))
+        x2 = self.proj_without_att(x)
+
+        return x1 + x2
+
+    def _project_master(self, x, master, att_map):
+
+        x1 = self.proj_with_attM(torch.matmul(
+            att_map.squeeze(-1).unsqueeze(1), x))
+        x2 = self.proj_without_attM(master)
+
+        return x1 + x2
+
+    def _apply_BN(self, x):
+        org_size = x.size()
+        x = x.view(-1, org_size[-1])
+        x = self.bn(x)
+        x = x.view(org_size)
+
+        return x
+
+    def _init_new_params(self, *size):
+        out = nn.Parameter(torch.FloatTensor(*size))
+        nn.init.xavier_normal_(out)
+        return out
+

audio/aasist3/model/kan.py

@@ -0,0 +1,213 @@
+import torch, math, torch.nn.functional as F
+
+
+class KANLinear(torch.nn.Module):
+    def __init__(
+        self,
+        in_features,
+        out_features,
+        grid_size=16,
+        spline_order=4,
+        scale_noise=0.1,
+        scale_base=1.0,
+        scale_spline=1.0,
+        enable_standalone_scale_spline=True,
+        base_activation=torch.nn.PReLU,
+        grid_eps=0.02,
+        grid_range=[-1, 1],
+    ):
+        super(KANLinear, self).__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.grid_size = grid_size
+        self.spline_order = spline_order
+
+        h = (grid_range[1] - grid_range[0]) / grid_size
+        grid = (
+            (
+                torch.arange(-spline_order, grid_size + spline_order + 1) * h
+                + grid_range[0]
+            )
+            .expand(in_features, -1)
+            .contiguous()
+        )
+        self.register_buffer("grid", grid)
+
+        self.base_weight = torch.nn.Parameter(torch.Tensor(out_features, in_features))
+        self.spline_weight = torch.nn.Parameter(
+            torch.Tensor(out_features, in_features, grid_size + spline_order)
+        )
+        if enable_standalone_scale_spline:
+            self.spline_scaler = torch.nn.Parameter(
+                torch.Tensor(out_features, in_features)
+            )
+
+        self.scale_noise = scale_noise
+        self.scale_base = scale_base
+        self.scale_spline = scale_spline
+        self.enable_standalone_scale_spline = enable_standalone_scale_spline
+        self.base_activation = base_activation()
+        self.grid_eps = grid_eps
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        torch.nn.init.kaiming_uniform_(self.base_weight, a=math.sqrt(5) * self.scale_base)
+        with torch.no_grad():
+            noise = (
+                (
+                    torch.rand(self.grid_size + 1, self.in_features, self.out_features)
+                    - 1 / 2
+                )
+                * self.scale_noise
+                / self.grid_size
+            )
+            self.spline_weight.data.copy_(
+                (self.scale_spline if not self.enable_standalone_scale_spline else 1.0)
+                * self.curve2coeff(
+                    self.grid.T[self.spline_order : -self.spline_order],
+                    noise,
+                )
+            )
+            if self.enable_standalone_scale_spline:
+                # torch.nn.init.constant_(self.spline_scaler, self.scale_spline)
+                torch.nn.init.kaiming_uniform_(self.spline_scaler, a=math.sqrt(5) * self.scale_spline)
+
+    def b_splines(self, x: torch.Tensor):
+        """
+        Compute the B-spline bases for the given input tensor.
+
+        Args:
+            x (torch.Tensor): Input tensor of shape (batch_size, in_features).
+
+        Returns:
+            torch.Tensor: B-spline bases tensor of shape (batch_size, in_features, grid_size + spline_order).
+        """
+        assert x.dim() == 2 and x.size(1) == self.in_features
+
+        grid: torch.Tensor = (
+            self.grid
+        )  # (in_features, grid_size + 2 * spline_order + 1)
+        x = x.unsqueeze(-1)
+        bases = ((x >= grid[:, :-1]) & (x < grid[:, 1:])).to(x.dtype)
+        for k in range(1, self.spline_order + 1):
+            bases = (
+                (x - grid[:, : -(k + 1)])
+                / (grid[:, k:-1] - grid[:, : -(k + 1)])
+                * bases[:, :, :-1]
+            ) + (
+                (grid[:, k + 1 :] - x)
+                / (grid[:, k + 1 :] - grid[:, 1:(-k)])
+                * bases[:, :, 1:]
+            )
+
+        assert bases.size() == (
+            x.size(0),
+            self.in_features,
+            self.grid_size + self.spline_order,
+        )
+        return bases.contiguous()
+
+    def curve2coeff(self, x: torch.Tensor, y: torch.Tensor):
+        """
+        Compute the coefficients of the curve that interpolates the given points.
+
+        Args:
+            x (torch.Tensor): Input tensor of shape (batch_size, in_features).
+            y (torch.Tensor): Output tensor of shape (batch_size, in_features, out_features).
+
+        Returns:
+            torch.Tensor: Coefficients tensor of shape (out_features, in_features, grid_size + spline_order).
+        """
+        assert x.dim() == 2 and x.size(1) == self.in_features
+        assert y.size() == (x.size(0), self.in_features, self.out_features)
+
+        A = self.b_splines(x).transpose(
+            0, 1
+        )  # (in_features, batch_size, grid_size + spline_order)
+        B = y.transpose(0, 1)  # (in_features, batch_size, out_features)
+        solution = torch.linalg.lstsq(
+            A, B
+        ).solution  # (in_features, grid_size + spline_order, out_features)
+        result = solution.permute(
+            2, 0, 1
+        )  # (out_features, in_features, grid_size + spline_order)
+
+        assert result.size() == (
+            self.out_features,
+            self.in_features,
+            self.grid_size + self.spline_order,
+        )
+        return result.contiguous()
+
+    @property
+    def scaled_spline_weight(self):
+        return self.spline_weight * (
+            self.spline_scaler.unsqueeze(-1)
+            if self.enable_standalone_scale_spline
+            else 1.0
+        )
+
+    def forward(self, x: torch.Tensor):
+        assert x.size(-1) == self.in_features
+        original_shape = x.shape
+        x = x.reshape(-1, self.in_features)
+
+        base_output = F.linear(self.base_activation(x), self.base_weight)
+        spline_output = F.linear(
+            self.b_splines(x).view(x.size(0), -1),
+            self.scaled_spline_weight.reshape(self.out_features, -1),
+        )
+        output = base_output + spline_output
+        # print(*original_shape[:-1], output.shape)
+        output = output.view(*original_shape[:-1], self.out_features)
+        return output
+
+    @torch.no_grad()
+    def update_grid(self, x: torch.Tensor, margin=0.01):
+        assert x.dim() == 2 and x.size(1) == self.in_features
+        batch = x.size(0)
+
+        splines = self.b_splines(x)  # (batch, in, coeff)
+        splines = splines.permute(1, 0, 2)  # (in, batch, coeff)
+        orig_coeff = self.scaled_spline_weight  # (out, in, coeff)
+        orig_coeff = orig_coeff.permute(1, 2, 0)  # (in, coeff, out)
+        unreduced_spline_output = torch.bmm(splines, orig_coeff)  # (in, batch, out)
+        unreduced_spline_output = unreduced_spline_output.permute(
+            1, 0, 2
+        )  # (batch, in, out)
+
+        # sort each channel individually to collect data distribution
+        x_sorted = torch.sort(x, dim=0)[0]
+        grid_adaptive = x_sorted[
+            torch.linspace(
+                0, batch - 1, self.grid_size + 1, dtype=torch.int64, device=x.device
+            )
+        ]
+
+        uniform_step = (x_sorted[-1] - x_sorted[0] + 2 * margin) / self.grid_size
+        grid_uniform = (
+            torch.arange(
+                self.grid_size + 1, dtype=torch.float32, device=x.device
+            ).unsqueeze(1)
+            * uniform_step
+            + x_sorted[0]
+            - margin
+        )
+
+        grid = self.grid_eps * grid_uniform + (1 - self.grid_eps) * grid_adaptive
+        grid = torch.concatenate(
+            [
+                grid[:1]
+                - uniform_step
+                * torch.arange(self.spline_order, 0, -1, device=x.device).unsqueeze(1),
+                grid,
+                grid[-1:]
+                + uniform_step
+                * torch.arange(1, self.spline_order + 1, device=x.device).unsqueeze(1),
+            ],
+            dim=0,
+        )
+
+        self.grid.copy_(grid.T)
+        self.spline_weight.data.copy_(self.curve2coeff(x, unreduced_spline_output))

audio/aasist3/model/pool.py

@@ -0,0 +1,45 @@
+import torch, torch.nn as nn
+from typing import Union
+
+from .kan import KANLinear
+
+
+class GraphPool(nn.Module):
+    def __init__(self, k: float, in_dim: int, p: Union[float, int], size):
+        super().__init__()
+        self.k = k
+        self.sigmoid = nn.Sigmoid()
+        self.proj = KANLinear(in_dim, 1)
+        self.drop = nn.Dropout(p=p) if p > 0 else nn.Identity()
+        self.in_dim = in_dim
+
+    def forward(self, h):
+        Z = self.drop(h)
+        weights = self.proj(Z)
+        scores = self.sigmoid(weights)
+        new_h = self.top_k_graph(scores, h, self.k)
+
+        return new_h
+
+    def top_k_graph(self, scores, h, k):
+        """
+        args
+        =====
+        scores: attention-based weights (#bs, #node, 1)
+        h: graph data (#bs, #node, #dim)
+        k: ratio of remaining nodes, (float)
+
+        returns
+        =====
+        h: graph pool applied data (#bs, #node', #dim)
+        """
+        _, n_nodes, n_feat = h.size()
+        n_nodes = max(int(n_nodes * k), 1)
+        _, idx = torch.topk(scores, n_nodes, dim=1)
+        idx = idx.expand(-1, -1, n_feat)
+
+        h = h * scores
+        h = torch.gather(h, 1, idx)
+
+        return h
+

audio/aasist3/model/residual.py

@@ -0,0 +1,56 @@
+import torch.nn as nn
+
+class Residual_block(nn.Module):
+    def __init__(self, nb_filts, first=False):
+        super().__init__()
+        self.first = first
+
+        if not self.first:
+            self.bn1 = nn.BatchNorm2d(num_features=nb_filts[0])
+        self.conv1 = nn.Conv2d(in_channels=nb_filts[0],
+                               out_channels=nb_filts[1],
+                               kernel_size=(2, 3),
+                               padding=(1, 1),
+                               stride=1)
+        self.selu = nn.SELU()
+
+        self.bn2 = nn.BatchNorm2d(num_features=nb_filts[1])
+        self.conv2 = nn.Conv2d(in_channels=nb_filts[1],
+                               out_channels=nb_filts[1],
+                               kernel_size=(2, 3),
+                               padding=(0, 1),
+                               stride=1)
+
+        if nb_filts[0] != nb_filts[1]:
+            self.downsample = True
+            self.conv_downsample = nn.Conv2d(in_channels=nb_filts[0],
+                                             out_channels=nb_filts[1],
+                                             padding=(0, 1),
+                                             kernel_size=(1, 3),
+                                             stride=1)
+
+        else:
+            self.downsample = False
+        # self.mp = nn.MaxPool2d((1, 3))  # self.mp = nn.MaxPool2d((1,4))
+
+    def forward(self, x):
+        identity = x
+        if not self.first:
+            out = self.bn1(x)
+            out = self.selu(out)
+        else:
+            out = x
+        out = self.conv1(x)
+
+        # print('out',out.shape)
+        out = self.bn2(out)
+        out = self.selu(out)
+        # print('out',out.shape)
+        out = self.conv2(out)
+        #print('conv2 out',out.shape)
+        if self.downsample:
+            identity = self.conv_downsample(identity)
+
+        out += identity
+        # out = self.mp(out)
+        return out

audio/aasist3/model/wav2vec.py

@@ -0,0 +1,82 @@
+import torch, torch.nn as nn
+from transformers import Wav2Vec2Model, Wav2Vec2Config
+
+
+class Wav2Vec2Encoder(nn.Module):
+    """SSL encoder based on Hugging Face's Wav2Vec2 model."""
+
+    def __init__(self,
+                 model_name_or_path: str = "facebook/wav2vec2-large-xlsr-53",
+                 ssl_out_dim: int = 768,
+                 use_ssl_n_layers: int = None,
+                 freeze_ssl_n_layers: int = 0,
+                 output_attentions: bool = False,
+                 output_hidden_states: bool = False,
+                 normalize_waveform: bool = True,
+                 cache_dir: str = "weights",
+                 load_pretrained: bool = True):
+        """Initialize the Wav2Vec2 encoder.
+
+        Args:
+            model_name_or_path: HuggingFace model name or path to local model.
+            ssl_out_dim: Output dimension of the Wav2Vec2 encoder.
+            use_ssl_n_layers: Number of Wav2Vec2 layers to use. If None, use all layers.
+            freeze_ssl_n_layers: Number of Wav2Vec2 layers to freeze during training.
+            output_attentions: Whether to output attentions.
+            output_hidden_states: Whether to output hidden states.
+            normalize_waveform: Whether to normalize the waveform input.
+            cache_dir: Directory to cache pretrained models.
+            load_pretrained: Whether to load pretrained weights. If False, initializes with random weights.
+        """
+        super().__init__()
+
+        self.model_name_or_path = model_name_or_path
+        self.ssl_out_dim = ssl_out_dim
+        self.use_ssl_n_layers = use_ssl_n_layers
+        self.freeze_ssl_n_layers = freeze_ssl_n_layers
+        self.output_attentions = output_attentions
+        self.output_hidden_states = output_hidden_states
+        self.normalize_waveform = normalize_waveform
+
+        if load_pretrained:
+            self.model = Wav2Vec2Model.from_pretrained(model_name_or_path, cache_dir=cache_dir)
+        else:
+            config = Wav2Vec2Config.from_pretrained(
+                model_name_or_path, 
+                cache_dir=cache_dir,
+                local_files_only=False
+            )
+            self.model = Wav2Vec2Model(config)
+            self.model.init_weights()
+
+    def forward(self, x):
+        """Forward pass through the Wav2Vec2 encoder.
+
+        Args:
+            x: Input tensor of shape (batch_size, sequence_length, channels)
+
+        Returns:
+            Extracted features of shape (batch_size, sequence_length, ssl_out_dim)
+        """
+        # Handle shape: convert (batch_size, sequence_length, channels) to (batch_size, sequence_length)
+        if x.ndim == 3:
+            x = x.squeeze(-1)  # Remove channel dimension if present
+
+        if self.normalize_waveform:
+            x = x / (torch.max(torch.abs(x), dim=1, keepdim=True)[0] + 1e-8)
+
+        outputs = self.model(
+            x,
+            output_attentions=self.output_attentions,
+            output_hidden_states=self.output_hidden_states,
+            return_dict=True
+        )
+
+        last_hidden_state = outputs.last_hidden_state
+
+        if self.use_ssl_n_layers is not None and self.output_hidden_states and outputs.hidden_states is not None:
+            selected = outputs.hidden_states[-self.use_ssl_n_layers:]
+            last_hidden_state = torch.mean(torch.stack(selected, dim=0), dim=0)
+        del outputs
+
+        return last_hidden_state

audio/aasist3/requirements.txt

@@ -0,0 +1,11 @@
+torch==2.5.1
+torchaudio==2.5.1
+transformers>=4.40.0
+huggingface-hub>=0.20.0
+safetensors>=0.4.0
+numpy<2.0
+soundfile
+fastapi
+uvicorn[standard]
+pydantic
+python-multipart

audio/nes2net/.gitignore

@@ -0,0 +1 @@
+weights/

audio/nes2net/Dockerfile

@@ -0,0 +1,41 @@
+FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
+
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PYTHONUNBUFFERED=1
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    python3 python3-pip python3-dev \
+    git ffmpeg libsndfile1 wget \
+    build-essential g++ \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN ln -sf /usr/bin/python3 /usr/bin/python
+
+WORKDIR /app
+
+# Install PyTorch with CUDA 12.1
+RUN pip install --no-cache-dir \
+    torch==2.5.1 torchaudio==2.5.1 \
+    --index-url https://download.pytorch.org/whl/cu121
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Clone fairseq with patched C extensions (same as ShiftySpeech)
+RUN git clone https://github.com/facebookresearch/fairseq.git /app/fairseq_repo && \
+    cd /app/fairseq_repo && \
+    git checkout a54021305d6b3c4c5959ac9395135f63202db8f1 && \
+    sed -i 's/ext_modules=extensions/ext_modules=[]/' setup.py && \
+    pip install --no-cache-dir --no-deps -e .
+
+COPY model_scripts /app/model_scripts
+COPY api.py .
+RUN mkdir -p /app/weights
+COPY weights/ /app/weights/
+
+EXPOSE 8004
+
+RUN adduser --disabled-password --gecos '' appuser
+USER appuser
+
+CMD ["python", "api.py"]

audio/nes2net/api.py

@@ -0,0 +1,315 @@
+"""Nes2Net (XLSR + Nested Res2Net TDNN) Audio Deepfake Detection API.
+
+Detects synthetic speech using the Nes2Net model architecture:
+- Frontend: XLSR wav2vec 2.0 (Self-Supervised Learning)
+- Backend: Nested Res2Net TDNN with SE modules
+
+Reference: https://github.com/TianchiLiu/Nes2Net
+"""
+
+import argparse
+import base64
+import io
+import logging
+import os
+import platform
+import sys
+import time
+import warnings
+from typing import Optional
+
+import librosa
+import numpy as np
+import torch
+import uvicorn
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field
+
+# Suppress deprecation warnings from fairseq/omegaconf compatibility
+warnings.filterwarnings("ignore", category=DeprecationWarning)
+
+# Monkey-patch omegaconf for fairseq compatibility (older fairseq
+# expects is_primitive_type which was removed in newer omegaconf).
+import omegaconf._utils as _omegaconf_utils
+
+if not hasattr(_omegaconf_utils, "is_primitive_type"):
+    _omegaconf_utils.is_primitive_type = lambda t: t in (int, float, bool, str, bytes)
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger("nes2net_api")
+
+# Add the model code to the path
+if "/app" not in sys.path:
+    sys.path.insert(0, "/app")
+
+# Import model class (deferred to allow path setup)
+try:
+    from model_scripts.wav2vec2_Nes2Net_X import (
+        wav2vec2_Nes2Net_no_Res_w_allT as Nes2NetModel,
+    )
+except ImportError as e:
+    logger.error(f"Failed to import Nes2Net model: {e}")
+    Nes2NetModel = None
+
+# Constants
+MODEL_NAME = "nes2net"
+MODEL_ID = "nes2net_xlsr_itw_valaug"
+WEIGHTS_PATH = "/app/weights/nes2net_itw_valaug.pt"
+
+
+def _get_device():
+    """Select optimal device: MPS (Apple) > CUDA (NVIDIA) > CPU."""
+    override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower()
+    if override == "cpu":
+        return torch.device("cpu")
+    if override == "cuda" and torch.cuda.is_available():
+        return torch.device("cuda")
+    if (
+        override == "mps"
+        and hasattr(torch.backends, "mps")
+        and torch.backends.mps.is_available()
+    ):
+        return torch.device("mps")
+    if override:
+        pass  # Invalid override, fall through to auto-detect
+    if (
+        platform.system() == "Darwin"
+        and hasattr(torch.backends, "mps")
+        and torch.backends.mps.is_available()
+    ):
+        return torch.device("mps")
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    return torch.device("cpu")
+
+
+DEVICE = _get_device()
+
+if DEVICE.type == "cuda":
+    torch.backends.cudnn.benchmark = True
+    torch.set_float32_matmul_precision("high")
+
+if DEVICE.type == "cuda":
+    logger.info(
+        "Device: cuda (%s, %.1f GB VRAM)",
+        torch.cuda.get_device_name(0),
+        torch.cuda.get_device_properties(0).total_memory / 1024**3,
+    )
+else:
+    logger.warning(
+        "Device: %s (no CUDA available -- check nvidia-container-toolkit)",
+        DEVICE,
+    )
+
+SAMPLE_RATE = 16000
+TARGET_SAMPLES = 64600  # ~4.04 seconds at 16kHz
+
+# Global model instance
+model = None
+
+
+class AudioInput(BaseModel):
+    """Request schema for audio deepfake detection."""
+
+    audio_data: str = Field(
+        ..., description="Base64 encoded audio string (WAV/MP3/etc)"
+    )
+    threshold: Optional[float] = Field(
+        0.5, ge=0.0, le=1.0, description="Classification threshold"
+    )
+
+
+app = FastAPI(
+    title="Nes2Net Audio Deepfake Detection API",
+    description=(
+        "Service for detecting synthetic speech using the "
+        "Nes2Net model (XLSR wav2vec 2.0 + Nested Res2Net TDNN)."
+    ),
+    version="1.0.0",
+)
+
+
+def load_model():
+    """Load the Nes2Net model with fine-tuned weights.
+
+    Returns:
+        The loaded model, or None if loading fails.
+    """
+    global model
+    if model is not None:
+        return model
+
+    logger.info(f"Loading Nes2Net model onto {DEVICE}...")
+
+    if Nes2NetModel is None:
+        logger.error("Nes2Net model class not available.")
+        return None
+
+    if not os.path.exists(WEIGHTS_PATH):
+        logger.error(f"Model weights not found at {WEIGHTS_PATH}")
+        return None
+
+    try:
+        args = argparse.Namespace(
+            n_output_logits=2,
+            dilation=2,
+            pool_func="mean",
+            SE_ratio=[1],
+            Nes_ratio=[8, 8],
+        )
+        model = Nes2NetModel(args, str(DEVICE))
+
+        # Load fine-tuned weights
+        try:
+            state_dict = torch.load(
+                WEIGHTS_PATH,
+                map_location=DEVICE,
+                weights_only=False,
+            )
+        except TypeError:
+            state_dict = torch.load(WEIGHTS_PATH, map_location=DEVICE)
+
+        model.load_state_dict(state_dict)
+        model.to(DEVICE)
+        model.eval()
+
+        logger.info("Nes2Net model loaded successfully.")
+        return model
+    except Exception as e:
+        logger.exception(f"Failed to load Nes2Net model: {e}")
+        model = None
+        return None
+
+
+@app.on_event("startup")
+async def startup_event():
+    """Load model on service startup."""
+    load_model()
+
+
+def _gpu_health_info() -> dict:
+    """Return GPU metrics for the health endpoint."""
+    if torch.cuda.is_available() and DEVICE.type == "cuda":
+        return {
+            "gpu_name": torch.cuda.get_device_name(0),
+            "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2),
+            "vram_total_mb": round(
+                torch.cuda.get_device_properties(0).total_memory / 1024**2
+            ),
+        }
+    return {}
+
+
+@app.get("/health")
+async def health():
+    """Health check endpoint."""
+    return {
+        "status": "healthy" if model is not None else "degraded",
+        "model": MODEL_NAME,
+        "model_id": MODEL_ID,
+        "device": str(DEVICE),
+        "weights_found": os.path.exists(WEIGHTS_PATH),
+        **_gpu_health_info(),
+    }
+
+
+def preprocess_audio(audio_bytes: bytes) -> torch.Tensor:
+    """Preprocess audio for Nes2Net inference.
+
+    Loads audio, resamples to 16kHz mono, and pads/trims
+    to TARGET_SAMPLES using tiling (matching original training
+    preprocessing).
+
+    Args:
+        audio_bytes: Raw audio file bytes.
+
+    Returns:
+        Audio tensor of shape (1, TARGET_SAMPLES).
+
+    Raises:
+        ValueError: If audio preprocessing fails.
+    """
+    try:
+        logger.info("Starting audio preprocessing...")
+        audio, sr = librosa.load(io.BytesIO(audio_bytes), sr=SAMPLE_RATE, mono=True)
+        logger.info(f"Audio loaded. Length: {len(audio)} samples at {sr}Hz")
+
+        # Pad/trim to TARGET_SAMPLES using tiling
+        if len(audio) >= TARGET_SAMPLES:
+            audio = audio[:TARGET_SAMPLES]
+        else:
+            num_repeats = TARGET_SAMPLES // len(audio) + 1
+            audio = np.tile(audio, num_repeats)[:TARGET_SAMPLES]
+
+        logger.info(f"Audio padded/trimmed to {TARGET_SAMPLES} samples")
+
+        audio_tensor = torch.FloatTensor(audio).unsqueeze(0).to(DEVICE)
+        return audio_tensor
+    except Exception as e:
+        logger.error(f"Error preprocessing audio: {e}")
+        raise ValueError(f"Audio preprocessing failed: {str(e)}")
+
+
+@app.post("/predict")
+async def predict(input_data: AudioInput):
+    """Run deepfake detection on base64-encoded audio.
+
+    The model outputs 2 logits: [spoof_score, bonafide_score].
+    Class 0 = spoof (fake), Class 1 = bonafide (real).
+    The returned probability is the spoof/fake probability.
+    """
+    if model is None:
+        if load_model() is None:
+            raise HTTPException(status_code=503, detail="Model not loaded")
+
+    try:
+        start_time = time.time()
+        logger.info(
+            f"Prediction request. Data size: " f"{len(input_data.audio_data)} chars"
+        )
+
+        # Decode base64 audio
+        audio_bytes = base64.b64decode(input_data.audio_data)
+
+        # Preprocess
+        audio_tensor = preprocess_audio(audio_bytes)
+
+        # Inference
+        logger.info("Starting model inference...")
+        with torch.no_grad():
+            output = model(audio_tensor)
+
+            # output shape: [batch, 2]
+            # Index 0 = spoof logit, Index 1 = bonafide logit
+            probs = torch.softmax(output, dim=1)
+            prob_fake = probs[0, 0].item()
+
+        prediction = 1 if prob_fake >= input_data.threshold else 0
+        verdict = "fake" if prediction == 1 else "real"
+        inference_time = time.time() - start_time
+
+        logger.info(
+            f"Prediction: {verdict} (prob_fake={prob_fake:.4f}, "
+            f"time={inference_time:.3f}s)"
+        )
+
+        return {
+            "model": MODEL_NAME,
+            "probability": float(prob_fake),
+            "prediction": int(prediction),
+            "class": verdict,
+            "inference_time": float(inference_time),
+        }
+
+    except Exception as e:
+        logger.exception(f"Error during prediction: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+if __name__ == "__main__":
+    port = int(os.environ.get("MODEL_PORT", 8004))
+    uvicorn.run(app, host="0.0.0.0", port=port)

audio/nes2net/model_scripts/wav2vec2_Nes2Net_X.py

@@ -0,0 +1,317 @@
+import math
+
+import fairseq
+import torch
+import torch.nn as nn
+
+___author__ = "Tianchi Liu"
+__email__ = "tianchi_liu@u.nus.edu"
+# modified from the model script from Hemlata Tak
+
+
+class SSLModel(nn.Module):
+    def __init__(self, device):
+        super(SSLModel, self).__init__()
+        cp_path = (
+            "/app/weights/xlsr2_300m.pt"  # Change the pre-trained XLSR model path.
+        )
+        model, cfg, task = fairseq.checkpoint_utils.load_model_ensemble_and_task(
+            [cp_path]
+        )
+        self.model = model[0]
+        self.device = device
+        self.out_dim = 1024
+        return
+
+    def extract_feat(self, input_data):
+        # put the model to GPU if it not there
+        if (
+            next(self.model.parameters()).device != input_data.device
+            or next(self.model.parameters()).dtype != input_data.dtype
+        ):
+            self.model.to(input_data.device, dtype=input_data.dtype)
+            self.model.train()
+        if True:
+            # input should be in shape (batch, length)
+            if input_data.ndim == 3:
+                input_tmp = input_data[:, :, 0]
+            else:
+                input_tmp = input_data
+            # [batch, length, dim]
+            emb = self.model(input_tmp, mask=False, features_only=True)["x"]
+        return emb
+
+
+class SEModule(nn.Module):
+    def __init__(self, channels, SE_ratio=8):
+        super(SEModule, self).__init__()
+        self.se = nn.Sequential(
+            nn.AdaptiveAvgPool1d(1),
+            nn.Conv1d(channels, channels // SE_ratio, kernel_size=1, padding=0),
+            nn.ReLU(),
+            nn.Conv1d(channels // SE_ratio, channels, kernel_size=1, padding=0),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, input):
+        x = self.se(input)
+        return input * x
+
+
+class Bottle2neck(nn.Module):
+
+    def __init__(
+        self, inplanes, planes, kernel_size=None, dilation=None, scale=8, SE_ratio=8
+    ):
+        super(Bottle2neck, self).__init__()
+        width = int(math.floor(planes / scale))
+        self.conv1 = nn.Conv1d(inplanes, width * scale, kernel_size=1)
+        self.bn1 = nn.BatchNorm1d(width * scale)
+        self.nums = scale - 1
+        convs = []
+        bns = []
+        weighted_sum = []
+        num_pad = math.floor(kernel_size / 2) * dilation
+        for i in range(self.nums):
+            convs.append(
+                nn.Conv2d(
+                    width,
+                    width,
+                    kernel_size=(kernel_size, 1),
+                    dilation=(dilation, 1),
+                    padding=(num_pad, 0),
+                )
+            )
+            bns.append(nn.BatchNorm2d(width))
+            initial_value = torch.ones(1, 1, 1, i + 2) * (1 / (i + 2))
+            weighted_sum.append(nn.Parameter(initial_value, requires_grad=True))
+        self.weighted_sum = nn.ParameterList(weighted_sum)
+        self.convs = nn.ModuleList(convs)
+        self.bns = nn.ModuleList(bns)
+        self.conv3 = nn.Conv1d(width * scale, planes, kernel_size=1)
+        self.bn3 = nn.BatchNorm1d(planes)
+        self.relu = nn.ReLU()
+        self.width = width
+        self.se = SEModule(planes, SE_ratio)
+
+    def forward(self, x):
+        residual = x
+        out = self.conv1(x)
+        out = self.relu(out)
+        out = self.bn1(out).unsqueeze(-1)  # bz c T 1
+
+        spx = torch.split(out, self.width, 1)
+        sp = spx[self.nums]
+        for i in range(self.nums):
+            sp = torch.cat((sp, spx[i]), -1)
+
+            sp = self.bns[i](self.relu(self.convs[i](sp)))
+            sp_s = sp * self.weighted_sum[i]
+            sp_s = torch.sum(sp_s, dim=-1, keepdim=False)
+
+            if i == 0:
+                out = sp_s
+            else:
+                out = torch.cat((out, sp_s), 1)
+        out = torch.cat((out, spx[self.nums].squeeze(-1)), 1)
+        out = self.conv3(out)
+        out = self.relu(out)
+        out = self.bn3(out)
+        out = self.se(out)
+        out += residual
+        return out
+
+
+class ASTP(nn.Module):
+    """Attentive statistics pooling: Channel- and context-dependent
+    statistics pooling, first used in ECAPA_TDNN.
+    """
+
+    def __init__(self, in_dim, bottleneck_dim=128, global_context_att=False):
+        super(ASTP, self).__init__()
+        self.global_context_att = global_context_att
+
+        # Use Conv1d with stride == 1 rather than Linear, then we don't
+        # need to transpose inputs.
+        if global_context_att:
+            self.linear1 = nn.Conv1d(
+                in_dim * 3, bottleneck_dim, kernel_size=1
+            )  # equals W and b in the paper
+        else:
+            self.linear1 = nn.Conv1d(
+                in_dim, bottleneck_dim, kernel_size=1
+            )  # equals W and b in the paper
+        self.linear2 = nn.Conv1d(
+            bottleneck_dim, in_dim, kernel_size=1
+        )  # equals V and k in the paper
+
+    def forward(self, x):
+        """
+        x: a 3-dimensional tensor in tdnn-based architecture (B,F,T)
+            or a 4-dimensional tensor in resnet architecture (B,C,F,T)
+            0-dim: batch-dimension, last-dim: time-dimension (frame-dimension)
+        """
+        if len(x.shape) == 4:
+            x = x.reshape(x.shape[0], x.shape[1] * x.shape[2], x.shape[3])
+        assert len(x.shape) == 3
+
+        if self.global_context_att:
+            context_mean = torch.mean(x, dim=-1, keepdim=True).expand_as(x)
+            context_std = torch.sqrt(
+                torch.var(x, dim=-1, keepdim=True) + 1e-10
+            ).expand_as(x)
+            x_in = torch.cat((x, context_mean, context_std), dim=1)
+        else:
+            x_in = x
+
+        # DON'T use ReLU here! ReLU may be hard to converge.
+        alpha = torch.tanh(self.linear1(x_in))  # alpha = F.relu(self.linear1(x_in))
+        alpha = torch.softmax(self.linear2(alpha), dim=2)
+        mean = torch.sum(alpha * x, dim=2)
+        var = torch.sum(alpha * (x**2), dim=2) - mean**2
+        std = torch.sqrt(var.clamp(min=1e-10))
+        return torch.cat([mean, std], dim=1)
+
+
+class Nested_Res2Net_TDNN(nn.Module):
+
+    def __init__(
+        self,
+        Nes_ratio=[8, 8],
+        input_channel=1024,
+        n_output_logits=2,
+        dilation=2,
+        pool_func="mean",
+        SE_ratio=[8],
+    ):
+
+        super(Nested_Res2Net_TDNN, self).__init__()
+        self.Nes_ratio = Nes_ratio[0]
+        assert input_channel % Nes_ratio[0] == 0
+        C = input_channel // Nes_ratio[0]
+        self.C = C
+        Build_in_Res2Nets = []
+        bns = []
+        for i in range(Nes_ratio[0] - 1):
+            Build_in_Res2Nets.append(
+                Bottle2neck(
+                    C,
+                    C,
+                    kernel_size=3,
+                    dilation=dilation,
+                    scale=Nes_ratio[1],
+                    SE_ratio=SE_ratio[0],
+                )
+            )
+            bns.append(nn.BatchNorm1d(C))
+        self.Build_in_Res2Nets = nn.ModuleList(Build_in_Res2Nets)
+        self.bns = nn.ModuleList(bns)
+        self.bn = nn.BatchNorm1d(1024)
+        self.relu = nn.ReLU()
+        self.pool_func = pool_func
+        if pool_func == "mean":
+            self.fc = nn.Linear(1024, n_output_logits)
+        elif pool_func == "ASTP":
+            self.pooling = ASTP(
+                in_dim=input_channel, bottleneck_dim=128, global_context_att=False
+            )
+            self.fc = nn.Linear(2048, n_output_logits)
+
+    def forward(self, x):
+        spx = torch.split(x, self.C, 1)
+        for i in range(self.Nes_ratio - 1):
+            if i == 0:
+                sp = spx[i]
+            else:
+                sp = sp + spx[i]
+            sp = self.Build_in_Res2Nets[i](sp)
+            sp = self.relu(sp)
+            sp = self.bns[i](sp)
+            if i == 0:
+                out = sp
+            else:
+                out = torch.cat((out, sp), 1)
+        out = torch.cat((out, spx[-1]), 1)
+        out = self.bn(out)
+        out = self.relu(out)
+        if self.pool_func == "mean":
+            out = torch.mean(out, dim=-1)
+        elif self.pool_func == "ASTP":
+            out = self.pooling(out)
+        out = self.fc(out)
+        return out
+
+
+class wav2vec2_Nes2Net_no_Res_w_allT(nn.Module):
+    def __init__(self, args, device):
+        super().__init__()
+        self.device = device
+
+        self.n_output_logits = args.n_output_logits
+
+        ####
+        # create network wav2vec 2.0
+        ####
+        self.ssl_model = SSLModel(self.device)
+        self.Nested_Res2Net_TDNN = Nested_Res2Net_TDNN(
+            Nes_ratio=args.Nes_ratio,
+            input_channel=1024,
+            n_output_logits=self.n_output_logits,
+            dilation=args.dilation,
+            pool_func=args.pool_func,
+            SE_ratio=args.SE_ratio,
+        )
+
+    def forward(self, x):
+        # -------pre-trained Wav2vec model fine tunning ------------------------##
+        x_ssl_feat = self.ssl_model.extract_feat(x.squeeze(-1))
+        x_ssl_feat = x_ssl_feat.permute(0, 2, 1)
+        output = self.Nested_Res2Net_TDNN(x_ssl_feat)
+
+        return output
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--n_output_logits", type=int, default=2)
+    parser.add_argument("--dilation", type=int, default=2)  # not important
+    parser.add_argument(
+        "--pool_func",
+        type=str,
+        default="mean",
+        choices=["mean", "ASTP"],
+        help="pooling function, choose from mean and ASTP",
+    )
+    parser.add_argument(
+        "--Nes_ratio",
+        type=int,
+        nargs="+",
+        default=[8, 8],
+        help="Nes_ratio, from outer to inner",
+    )
+    parser.add_argument(
+        "--SE_ratio",
+        type=int,
+        nargs="+",
+        default=[1],
+        help="SE downsampling ratio in the bottleneck",
+    )
+    args = parser.parse_args()
+
+    model = wav2vec2_Nes2Net_no_Res_w_allT(args=args, device="cpu")
+    x = torch.rand((4, 32000)).to("cpu")
+    model = model.to("cpu")
+    y = model(x)
+    print(y)
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print("all:", trainable_params)
+    trainable_params = sum(
+        p.numel() for p in model.ssl_model.parameters() if p.requires_grad
+    )
+    print("SSL:", trainable_params)
+    trainable_params = sum(
+        p.numel() for p in model.Nested_Res2Net_TDNN.parameters() if p.requires_grad
+    )
+    print("Backend:", trainable_params)

audio/nes2net/requirements.txt

@@ -0,0 +1,13 @@
+torch==2.5.1
+torchaudio==2.5.1
+numpy==1.23.5
+librosa==0.9.1
+soundfile
+scipy
+omegaconf
+hydra-core
+bitarray
+fastapi
+uvicorn[standard]
+pydantic
+python-multipart

audio/safeear/Dockerfile

@@ -0,0 +1,49 @@
+FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
+
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PYTHONUNBUFFERED=1
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    python3 python3-pip python3-dev \
+    git ffmpeg libsndfile1 wget \
+    build-essential \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN ln -sf /usr/bin/python3 /usr/bin/python
+
+WORKDIR /app
+
+# Install PyTorch with CUDA 12.1
+RUN pip install --no-cache-dir \
+    torch==2.5.1 torchaudio==2.5.1 \
+    --index-url https://download.pytorch.org/whl/cu121
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Clone SafeEar repository (for model code imports)
+RUN git clone --depth 1 https://github.com/LetterLiGo/SafeEar.git /app/safeear_repo
+
+# Install the fairseq fork with C extensions PATCHED OUT
+# (same proven patch used by ShiftySpeech and Nes2Net --
+#  C extensions are not needed for checkpoint loading)
+WORKDIR /app/safeear_repo/fairseq_ours
+RUN sed -i 's/ext_modules=extensions/ext_modules=[]/' setup.py && \
+    pip install --no-cache-dir --no-deps -e .
+WORKDIR /app
+
+# Download model weights from HuggingFace
+RUN mkdir -p /app/weights && \
+    wget -q -O /app/weights/SpeechTokenizer.pt \
+      "https://huggingface.co/TEC2004/SafeEar-ASV19-spoof-detection/resolve/main/SpeechTokenizer.pt" && \
+    wget -q -O /app/weights/model.ckpt \
+      "https://huggingface.co/TEC2004/SafeEar-ASV19-spoof-detection/resolve/main/model.ckpt"
+
+COPY api.py .
+
+EXPOSE 8002
+
+RUN adduser --disabled-password --gecos '' appuser
+USER appuser
+
+CMD ["python", "api.py"]

audio/safeear/api.py

@@ -0,0 +1,321 @@
+"""SafeEar audio deepfake detection API service.
+
+Uses the SafeEar content privacy-preserving model (CCS 2024) to detect
+synthetic speech. Two-stage pipeline:
+  1. SpeechTokenizer (neural audio codec) decouples acoustic features
+  2. SafeEar1s (transformer classifier) detects spoofing from acoustic tokens
+
+Weights: HuggingFace TEC2004/SafeEar-ASV19-spoof-detection
+"""
+
+import base64
+import logging
+import os
+import sys
+import tempfile
+import time
+from typing import Optional
+
+import uvicorn
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger("safeear_api")
+
+import platform
+
+import librosa
+import numpy as np
+import torch
+
+# Add SafeEar repo to path for model imports
+SAFEEAR_REPO_PATH = os.environ.get(
+    "SAFEEAR_REPO_PATH",
+    os.path.join(os.path.dirname(__file__), "safeear_repo"),
+)
+if SAFEEAR_REPO_PATH not in sys.path:
+    sys.path.insert(0, SAFEEAR_REPO_PATH)
+
+
+def _get_device():
+    """Select optimal device: MPS (Apple) > CUDA (NVIDIA) > CPU."""
+    override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower()
+    if override == "cpu":
+        return torch.device("cpu")
+    if override == "cuda" and torch.cuda.is_available():
+        return torch.device("cuda")
+    if (
+        override == "mps"
+        and hasattr(torch.backends, "mps")
+        and torch.backends.mps.is_available()
+    ):
+        return torch.device("mps")
+    if (
+        platform.system() == "Darwin"
+        and hasattr(torch.backends, "mps")
+        and torch.backends.mps.is_available()
+    ):
+        return torch.device("mps")
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    return torch.device("cpu")
+
+
+# Constants
+MODEL_NAME = "safeear"
+WEIGHTS_DIR = os.environ.get(
+    "WEIGHTS_DIR",
+    os.path.join(os.path.dirname(__file__), "weights"),
+)
+DEVICE = _get_device()
+
+if DEVICE.type == "cuda":
+    torch.backends.cudnn.benchmark = True
+    torch.set_float32_matmul_precision("high")
+
+if DEVICE.type == "cuda":
+    logger.info(
+        "Device: cuda (%s, %.1f GB VRAM)",
+        torch.cuda.get_device_name(0),
+        torch.cuda.get_device_properties(0).total_memory / 1024**3,
+    )
+else:
+    logger.warning(
+        "Device: %s (no CUDA available -- check nvidia-container-toolkit)",
+        DEVICE,
+    )
+
+SAMPLE_RATE = 16000
+MAX_AUDIO_LENGTH = 64600  # ~4 seconds at 16kHz (ASVspoof standard)
+SOFTMAX_TEMPERATURE = 5.0  # Calibration temperature for out-of-distribution data
+NUM_INFERENCE_PASSES = 5  # Monte Carlo passes for stable predictions
+
+# Global model instances
+decouple_model = None
+detect_model = None
+
+
+class AudioInput(BaseModel):
+    """Schema for audio prediction requests."""
+
+    audio_data: str = Field(
+        ..., description="Base64 encoded audio string (WAV/MP3/etc)"
+    )
+    threshold: Optional[float] = Field(
+        0.5, ge=0.0, le=1.0, description="Classification threshold"
+    )
+
+
+app = FastAPI(
+    title="SafeEar Audio Deepfake Detection API",
+    description="Content privacy-preserving deepfake detection using SafeEar.",
+    version="1.0.0",
+)
+
+
+def load_models():
+    """Load both the decouple model (SpeechTokenizer) and detect model."""
+    global decouple_model, detect_model
+
+    if decouple_model is not None and detect_model is not None:
+        return True
+
+    speech_tokenizer_path = os.path.join(WEIGHTS_DIR, "SpeechTokenizer.pt")
+    checkpoint_path = os.path.join(WEIGHTS_DIR, "model.ckpt")
+
+    if not os.path.exists(speech_tokenizer_path):
+        logger.error(f"SpeechTokenizer weights not found: {speech_tokenizer_path}")
+        return False
+    if not os.path.exists(checkpoint_path):
+        logger.error(f"Model checkpoint not found: {checkpoint_path}")
+        return False
+
+    try:
+        # --- Load SpeechTokenizer (decouple model) ---
+        from safeear.models.decouple import SpeechTokenizer
+
+        logger.info("Loading SpeechTokenizer...")
+        decouple_model = SpeechTokenizer(
+            n_filters=64,
+            strides=[8, 5, 4, 2],
+            dimension=1024,
+            semantic_dimension=768,
+            bidirectional=True,
+            dilation_base=2,
+            residual_kernel_size=3,
+            n_residual_layers=1,
+            lstm_layers=2,
+            activation="ELU",
+            codebook_size=1024,
+            n_q=8,
+            sample_rate=16000,
+        )
+        st_state = torch.load(speech_tokenizer_path, map_location="cpu")
+        decouple_model.load_state_dict(st_state)
+        decouple_model.to(DEVICE)
+        decouple_model.eval()
+        logger.info("SpeechTokenizer loaded.")
+
+        # --- Load SafeEar1s (detect model) from Lightning checkpoint ---
+        from safeear.models.safeear import SafeEar1s, SE_Rawformer_front
+
+        logger.info("Loading SafeEar1s detect model...")
+        detect_model = SafeEar1s(
+            front=SE_Rawformer_front(),
+            embedding_dim=1024,
+            dropout_rate=0.1,
+            attention_dropout=0.1,
+            stochastic_depth=0.1,
+            num_layers=2,
+            num_heads=8,
+            num_classes=2,
+            positional_embedding="sine",
+            mlp_ratio=1.0,
+        )
+
+        # The .ckpt is a PyTorch Lightning checkpoint
+        ckpt = torch.load(checkpoint_path, map_location="cpu")
+        state_dict = ckpt.get("state_dict", ckpt)
+
+        # Lightning prefixes keys with "detect_model."
+        detect_state = {}
+        for k, v in state_dict.items():
+            if k.startswith("detect_model."):
+                detect_state[k.replace("detect_model.", "", 1)] = v
+
+        detect_model.load_state_dict(detect_state)
+        detect_model.to(DEVICE)
+        detect_model.eval()
+        logger.info("SafeEar1s detect model loaded.")
+        return True
+
+    except Exception as e:
+        logger.exception(f"Failed to load SafeEar models: {e}")
+        decouple_model = None
+        detect_model = None
+        return False
+
+
+def preprocess_audio(audio_bytes: bytes) -> torch.Tensor:
+    """Load audio bytes, resample to 16kHz mono, pad/trim."""
+    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
+        tmp.write(audio_bytes)
+        tmp_path = tmp.name
+
+    try:
+        waveform, _ = librosa.load(tmp_path, sr=SAMPLE_RATE, mono=True)
+    finally:
+        os.unlink(tmp_path)
+
+    if len(waveform) < MAX_AUDIO_LENGTH:
+        waveform = np.pad(waveform, (0, MAX_AUDIO_LENGTH - len(waveform)))
+    else:
+        waveform = waveform[:MAX_AUDIO_LENGTH]
+
+    # Shape: (1, 1, samples) -- batch=1, channels=1, time
+    tensor = torch.FloatTensor(waveform).unsqueeze(0).unsqueeze(0).to(DEVICE)
+    return tensor
+
+
+@app.on_event("startup")
+async def startup_event():
+    """Attempt to load models at startup."""
+    load_models()
+
+
+def _gpu_health_info() -> dict:
+    """Return GPU metrics for the health endpoint."""
+    if torch.cuda.is_available() and DEVICE.type == "cuda":
+        return {
+            "gpu_name": torch.cuda.get_device_name(0),
+            "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2),
+            "vram_total_mb": round(
+                torch.cuda.get_device_properties(0).total_memory / 1024**2
+            ),
+        }
+    return {}
+
+
+@app.get("/health")
+async def health():
+    """Return service health status and model availability."""
+    models_loaded = decouple_model is not None and detect_model is not None
+    return {
+        "status": "healthy" if models_loaded else "degraded",
+        "model": MODEL_NAME,
+        "device": str(DEVICE),
+        "weights_found": (
+            os.path.exists(os.path.join(WEIGHTS_DIR, "SpeechTokenizer.pt"))
+            and os.path.exists(os.path.join(WEIGHTS_DIR, "model.ckpt"))
+        ),
+        **_gpu_health_info(),
+    }
+
+
+@app.post("/predict")
+async def predict(input_data: AudioInput):
+    """Run SafeEar inference on base64-encoded audio data."""
+    if decouple_model is None or detect_model is None:
+        if not load_models():
+            raise HTTPException(status_code=503, detail="Models not loaded")
+
+    try:
+        start_time = time.time()
+        logger.info(
+            "Received prediction request. "
+            f"Data size: {len(input_data.audio_data)} chars"
+        )
+
+        audio_bytes = base64.b64decode(input_data.audio_data)
+        x_wav = preprocess_audio(audio_bytes)
+
+        with torch.no_grad():
+            # Step 1: Extract acoustic tokens via SpeechTokenizer
+            # forward() returns:
+            #   (reconstructed, commit_loss, semantic_feature, acoustic_tokens)
+            # layers=[0,1,2,3,4,5,6,7] means layer 0 goes to
+            # semantic_feature; layers 1-7 go to acoustic_tokens list
+            _, _, _, acoustic_tokens = decouple_model(
+                x_wav, layers=[0, 1, 2, 3, 4, 5, 6, 7]
+            )
+
+            # Step 2: Run detection model with Monte Carlo averaging
+            # SafeEar1s uses torch.randperm() in forward, so we average
+            # multiple passes for stable predictions
+            logit_sum = torch.zeros(1, 2, device=DEVICE)
+            for _ in range(NUM_INFERENCE_PASSES):
+                raw_logits, _ = detect_model(acoustic_tokens)
+                logit_sum += raw_logits
+            avg_logits = logit_sum / NUM_INFERENCE_PASSES
+
+            # Step 3: Get fake probability with temperature-scaled softmax
+            # The model produces extreme logits that saturate standard
+            # softmax. Temperature scaling preserves discrimination while
+            # giving more interpretable probabilities.
+            probs = torch.softmax(avg_logits / SOFTMAX_TEMPERATURE, dim=-1)
+            prob_fake = probs[0, 1].item()
+
+        prediction = 1 if prob_fake >= input_data.threshold else 0
+        verdict = "fake" if prediction == 1 else "real"
+        inference_time = time.time() - start_time
+
+        return {
+            "model": MODEL_NAME,
+            "probability": float(prob_fake),
+            "prediction": int(prediction),
+            "class": verdict,
+            "inference_time": float(inference_time),
+        }
+
+    except Exception as e:
+        logger.exception(f"Error during prediction: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+if __name__ == "__main__":
+    port = int(os.environ.get("MODEL_PORT", 8002))
+    uvicorn.run(app, host="0.0.0.0", port=port)

audio/safeear/download_weights.sh

@@ -0,0 +1,23 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+WEIGHTS_DIR="${WEIGHTS_DIR:-/app/weights}"
+REPO_DIR="${REPO_DIR:-/app/safeear_repo}"
+
+mkdir -p "$WEIGHTS_DIR"
+
+echo "==> Cloning SafeEar source repository..."
+if [ ! -d "$REPO_DIR/.git" ]; then
+    git clone --depth 1 https://github.com/LetterLiGo/SafeEar.git "$REPO_DIR"
+fi
+
+echo "==> Downloading SpeechTokenizer.pt from HuggingFace..."
+wget -q --show-progress -O "$WEIGHTS_DIR/SpeechTokenizer.pt" \
+  "https://huggingface.co/TEC2004/SafeEar-ASV19-spoof-detection/resolve/main/SpeechTokenizer.pt"
+
+echo "==> Downloading model.ckpt from HuggingFace..."
+wget -q --show-progress -O "$WEIGHTS_DIR/model.ckpt" \
+  "https://huggingface.co/TEC2004/SafeEar-ASV19-spoof-detection/resolve/main/model.ckpt"
+
+echo "==> Weights downloaded to $WEIGHTS_DIR"
+ls -lh "$WEIGHTS_DIR"

audio/safeear/requirements.txt

@@ -0,0 +1,15 @@
+torch==2.5.1
+torchaudio==2.5.1
+librosa>=0.10.0
+soundfile>=0.11.0
+numpy>=1.23.0
+einops>=0.7.0
+timm>=0.9.0
+hydra-core>=1.0.7
+omegaconf>=2.1.0
+pytorch-lightning>=1.6.0
+scipy>=1.11.0
+fastapi>=0.100.0
+uvicorn[standard]>=0.20.0
+python-multipart>=0.0.5
+pydantic>=2.0.0

audio/shiftyspeech/Dockerfile

@@ -0,0 +1,43 @@
+FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
+
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PYTHONUNBUFFERED=1
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    python3 python3-pip python3-dev \
+    git ffmpeg libsndfile1 \
+    build-essential g++ \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN ln -sf /usr/bin/python3 /usr/bin/python
+
+WORKDIR /app
+
+# Install PyTorch with CUDA 12.1
+RUN pip install --no-cache-dir \
+    torch==2.5.1 torchaudio==2.5.1 \
+    --index-url https://download.pytorch.org/whl/cu121
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Clone fairseq with patched C extensions
+RUN git clone https://github.com/facebookresearch/fairseq.git /app/fairseq_repo && \
+    cd /app/fairseq_repo && \
+    git checkout a54021305d6b3c4c5959ac9395135f63202db8f1 && \
+    sed -i 's/ext_modules=extensions/ext_modules=[]/' setup.py && \
+    pip install --no-cache-dir --no-deps -e .
+
+COPY synthetic_speech_detection /app/synthetic_speech_detection
+RUN mkdir -p /app/models
+COPY api.py .
+RUN mkdir -p /app/weights
+COPY weights/ /app/weights/
+RUN ln -sf /app/weights/xlsr_53_56k.pt /app/models/xlsr2_300m.pt
+
+EXPOSE 8001
+
+RUN adduser --disabled-password --gecos '' appuser
+USER appuser
+
+CMD ["python", "api.py"]

audio/shiftyspeech/api.py

@@ -0,0 +1,315 @@
+"""ShiftySpeech (SSL-AASIST) Audio Deepfake Detection API.
+
+Detects synthetic speech using the SSL-AASIST model architecture:
+- Frontend: XLSR wav2vec 2.0 (Self-Supervised Learning)
+- Backend: AASIST (Audio Anti-Spoofing using Integrated
+  Spectro-Temporal Graph Attention Networks)
+
+Reference: https://github.com/Ashigarg123/ShiftySpeech
+"""
+
+import base64
+import io
+import logging
+import os
+import platform
+import sys
+import time
+import warnings
+from typing import Optional
+
+import librosa
+import numpy as np
+import torch
+import uvicorn
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field
+
+# Suppress deprecation warnings from fairseq/omegaconf compatibility
+warnings.filterwarnings("ignore", category=DeprecationWarning)
+
+# Monkey-patch omegaconf for fairseq compatibility (older fairseq
+# expects is_primitive_type which was removed in newer omegaconf).
+import omegaconf._utils as _omegaconf_utils
+
+if not hasattr(_omegaconf_utils, "is_primitive_type"):
+    _omegaconf_utils.is_primitive_type = lambda t: t in (int, float, bool, str, bytes)
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger("shiftyspeech_api")
+
+# Add the SSL_Anti-spoofing model code to the path
+MODEL_CODE_PATH = "/app/synthetic_speech_detection/SSL_Anti-spoofing"
+if MODEL_CODE_PATH not in sys.path:
+    sys.path.insert(0, MODEL_CODE_PATH)
+
+# Import model class (deferred to allow path setup)
+try:
+    from model import Model as SSLAASISTModel
+except ImportError as e:
+    logger.error(f"Failed to import SSL-AASIST model: {e}")
+    SSLAASISTModel = None
+
+# Constants
+MODEL_NAME = "shiftyspeech"
+MODEL_ID = "ssl_aasist_augmented"
+WEIGHTS_PATH = "/app/weights/hfg_aug_1_2.pt"
+XLSR_DIR = "/app/models"
+
+
+def _get_device():
+    """Select optimal device: MPS (Apple) > CUDA (NVIDIA) > CPU."""
+    override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower()
+    if override == "cpu":
+        return torch.device("cpu")
+    if override == "cuda" and torch.cuda.is_available():
+        return torch.device("cuda")
+    if (
+        override == "mps"
+        and hasattr(torch.backends, "mps")
+        and torch.backends.mps.is_available()
+    ):
+        return torch.device("mps")
+    if override:
+        pass  # Invalid override, fall through to auto-detect
+    if (
+        platform.system() == "Darwin"
+        and hasattr(torch.backends, "mps")
+        and torch.backends.mps.is_available()
+    ):
+        return torch.device("mps")
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    return torch.device("cpu")
+
+
+DEVICE = _get_device()
+
+if DEVICE.type == "cuda":
+    torch.backends.cudnn.benchmark = True
+    torch.set_float32_matmul_precision("high")
+
+if DEVICE.type == "cuda":
+    logger.info(
+        "Device: cuda (%s, %.1f GB VRAM)",
+        torch.cuda.get_device_name(0),
+        torch.cuda.get_device_properties(0).total_memory / 1024**3,
+    )
+else:
+    logger.warning(
+        "Device: %s (no CUDA available -- check nvidia-container-toolkit)",
+        DEVICE,
+    )
+
+SAMPLE_RATE = 16000
+TARGET_SAMPLES = 64600  # ~4.04 seconds at 16kHz
+
+# Global model instance
+model = None
+
+
+class AudioInput(BaseModel):
+    """Request schema for audio deepfake detection."""
+
+    audio_data: str = Field(
+        ..., description="Base64 encoded audio string (WAV/MP3/etc)"
+    )
+    threshold: Optional[float] = Field(
+        0.5, ge=0.0, le=1.0, description="Classification threshold"
+    )
+
+
+app = FastAPI(
+    title="ShiftySpeech Audio Deepfake Detection API",
+    description=(
+        "Service for detecting synthetic speech using the "
+        "SSL-AASIST model (XLSR wav2vec 2.0 + AASIST backend)."
+    ),
+    version="1.0.0",
+)
+
+
+def load_model():
+    """Load the SSL-AASIST model with augmented weights.
+
+    Returns:
+        The loaded model, or None if loading fails.
+    """
+    global model
+    if model is not None:
+        return model
+
+    logger.info(f"Loading SSL-AASIST model onto {DEVICE}...")
+
+    if SSLAASISTModel is None:
+        logger.error("SSL-AASIST model class not available.")
+        return None
+
+    if not os.path.exists(WEIGHTS_PATH):
+        logger.error(f"Model weights not found at {WEIGHTS_PATH}")
+        return None
+
+    try:
+        # Ensure XLSR model directory exists for architecture init
+        os.makedirs(XLSR_DIR, exist_ok=True)
+
+        import argparse
+
+        args = argparse.Namespace()
+        model = SSLAASISTModel(args, str(DEVICE))
+
+        # Load fine-tuned weights (includes XLSR weights)
+        try:
+            state_dict = torch.load(
+                WEIGHTS_PATH,
+                map_location=DEVICE,
+                weights_only=False,
+            )
+        except TypeError:
+            state_dict = torch.load(WEIGHTS_PATH, map_location=DEVICE)
+
+        model.load_state_dict(state_dict)
+        model.to(DEVICE)
+        model.eval()
+
+        logger.info("SSL-AASIST model loaded successfully.")
+        return model
+    except Exception as e:
+        logger.exception(f"Failed to load SSL-AASIST model: {e}")
+        model = None
+        return None
+
+
+@app.on_event("startup")
+async def startup_event():
+    """Load model on service startup."""
+    load_model()
+
+
+def _gpu_health_info() -> dict:
+    """Return GPU metrics for the health endpoint."""
+    if torch.cuda.is_available() and DEVICE.type == "cuda":
+        return {
+            "gpu_name": torch.cuda.get_device_name(0),
+            "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2),
+            "vram_total_mb": round(
+                torch.cuda.get_device_properties(0).total_memory / 1024**2
+            ),
+        }
+    return {}
+
+
+@app.get("/health")
+async def health():
+    """Health check endpoint."""
+    return {
+        "status": "healthy" if model is not None else "degraded",
+        "model": MODEL_NAME,
+        "model_id": MODEL_ID,
+        "device": str(DEVICE),
+        "weights_found": os.path.exists(WEIGHTS_PATH),
+        **_gpu_health_info(),
+    }
+
+
+def preprocess_audio(audio_bytes: bytes) -> torch.Tensor:
+    """Preprocess audio for SSL-AASIST inference.
+
+    Loads audio, resamples to 16kHz mono, and pads/trims
+    to TARGET_SAMPLES using tiling (matching original training
+    preprocessing from data_utils.py).
+
+    Args:
+        audio_bytes: Raw audio file bytes.
+
+    Returns:
+        Audio tensor of shape (1, TARGET_SAMPLES).
+
+    Raises:
+        ValueError: If audio preprocessing fails.
+    """
+    try:
+        logger.info("Starting audio preprocessing...")
+        audio, sr = librosa.load(io.BytesIO(audio_bytes), sr=SAMPLE_RATE, mono=True)
+        logger.info(f"Audio loaded. Length: {len(audio)} samples at {sr}Hz")
+
+        # Pad/trim to TARGET_SAMPLES using tiling
+        # (matches original data_utils.pad function)
+        if len(audio) >= TARGET_SAMPLES:
+            audio = audio[:TARGET_SAMPLES]
+        else:
+            num_repeats = TARGET_SAMPLES // len(audio) + 1
+            audio = np.tile(audio, num_repeats)[:TARGET_SAMPLES]
+
+        logger.info(f"Audio padded/trimmed to {TARGET_SAMPLES} samples")
+
+        audio_tensor = torch.FloatTensor(audio).unsqueeze(0).to(DEVICE)
+        return audio_tensor
+    except Exception as e:
+        logger.error(f"Error preprocessing audio: {e}")
+        raise ValueError(f"Audio preprocessing failed: {str(e)}")
+
+
+@app.post("/predict")
+async def predict(input_data: AudioInput):
+    """Run deepfake detection on base64-encoded audio.
+
+    The model outputs 2 logits: [spoof_score, bonafide_score].
+    Class 0 = spoof (fake), Class 1 = bonafide (real).
+    The returned probability is the spoof/fake probability.
+    """
+    if model is None:
+        if load_model() is None:
+            raise HTTPException(status_code=503, detail="Model not loaded")
+
+    try:
+        start_time = time.time()
+        logger.info(
+            f"Prediction request. Data size: " f"{len(input_data.audio_data)} chars"
+        )
+
+        # Decode base64 audio
+        audio_bytes = base64.b64decode(input_data.audio_data)
+
+        # Preprocess
+        audio_tensor = preprocess_audio(audio_bytes)
+
+        # Inference
+        logger.info("Starting model inference...")
+        with torch.no_grad():
+            output = model(audio_tensor)
+
+            # output shape: [batch, 2]
+            # Index 0 = spoof logit, Index 1 = bonafide logit
+            probs = torch.softmax(output, dim=1)
+            prob_fake = probs[0, 0].item()
+
+        prediction = 1 if prob_fake >= input_data.threshold else 0
+        verdict = "fake" if prediction == 1 else "real"
+        inference_time = time.time() - start_time
+
+        logger.info(
+            f"Prediction: {verdict} (prob_fake={prob_fake:.4f}, "
+            f"time={inference_time:.3f}s)"
+        )
+
+        return {
+            "model": MODEL_NAME,
+            "probability": float(prob_fake),
+            "prediction": int(prediction),
+            "class": verdict,
+            "inference_time": float(inference_time),
+        }
+
+    except Exception as e:
+        logger.exception(f"Error during prediction: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+if __name__ == "__main__":
+    port = int(os.environ.get("MODEL_PORT", 8001))
+    uvicorn.run(app, host="0.0.0.0", port=port)

audio/shiftyspeech/evaluate.py

@@ -0,0 +1,210 @@
+"""Evaluate ShiftySpeech SSL-AASIST on the DeepSafe audio dataset.
+
+Reports accuracy, precision, recall, F1, EER, and per-file results.
+"""
+
+import os
+import sys
+import time
+import warnings
+
+warnings.filterwarnings("ignore", category=DeprecationWarning)
+
+# Monkey-patch omegaconf for fairseq compatibility
+import omegaconf._utils as _omegaconf_utils
+
+if not hasattr(_omegaconf_utils, "is_primitive_type"):
+    _omegaconf_utils.is_primitive_type = lambda t: t in (int, float, bool, str, bytes)
+
+import argparse
+
+import librosa
+import numpy as np
+import torch
+
+# Add model code to path
+SERVICE_DIR = os.path.dirname(os.path.abspath(__file__))
+MODEL_CODE_PATH = os.path.join(
+    SERVICE_DIR, "synthetic_speech_detection", "SSL_Anti-spoofing"
+)
+sys.path.insert(0, MODEL_CODE_PATH)
+
+from model import Model as SSLAASISTModel
+
+SAMPLE_RATE = 16000
+TARGET_SAMPLES = 64600
+DATASET_DIR = os.path.join(
+    SERVICE_DIR, os.pardir, os.pardir, os.pardir, "dataset", "audio"
+)
+DATASET_DIR = os.path.normpath(DATASET_DIR)
+
+
+def pad_audio(audio, target=TARGET_SAMPLES):
+    """Pad/trim audio to target length using tiling."""
+    if len(audio) >= target:
+        return audio[:target]
+    num_repeats = target // len(audio) + 1
+    return np.tile(audio, num_repeats)[:target]
+
+
+def compute_eer(target_scores, nontarget_scores):
+    """Compute Equal Error Rate."""
+    n_scores = target_scores.size + nontarget_scores.size
+    all_scores = np.concatenate((target_scores, nontarget_scores))
+    labels = np.concatenate(
+        (np.ones(target_scores.size), np.zeros(nontarget_scores.size))
+    )
+    indices = np.argsort(all_scores, kind="mergesort")
+    labels = labels[indices]
+    tar_trial_sums = np.cumsum(labels)
+    nontarget_trial_sums = nontarget_scores.size - (
+        np.arange(1, n_scores + 1) - tar_trial_sums
+    )
+    frr = np.concatenate((np.atleast_1d(0), tar_trial_sums / target_scores.size))
+    far = np.concatenate(
+        (
+            np.atleast_1d(1),
+            nontarget_trial_sums / nontarget_scores.size,
+        )
+    )
+    abs_diffs = np.abs(frr - far)
+    min_index = np.argmin(abs_diffs)
+    eer = np.mean((frr[min_index], far[min_index]))
+    return eer
+
+
+def main():
+    weights_path = os.path.join(SERVICE_DIR, "weights", "hfg_aug_1_2.pt")
+    os.makedirs(os.path.join(SERVICE_DIR, "models"), exist_ok=True)
+
+    print("=" * 70)
+    print("ShiftySpeech (SSL-AASIST) - DeepSafe Dataset Evaluation")
+    print("=" * 70)
+    print(f"Weights: {weights_path}")
+    print(f"Dataset: {DATASET_DIR}")
+    print(f"Device: cpu")
+    print()
+
+    # Load model
+    print("Loading model...")
+    start = time.time()
+    args_ns = argparse.Namespace()
+    ssl_model = SSLAASISTModel(args_ns, "cpu")
+    state_dict = torch.load(weights_path, map_location="cpu", weights_only=False)
+    ssl_model.load_state_dict(state_dict)
+    ssl_model.eval()
+    print(f"Model loaded in {time.time() - start:.1f}s")
+    print()
+
+    # Collect audio files
+    real_dir = os.path.join(DATASET_DIR, "real")
+    fake_dir = os.path.join(DATASET_DIR, "fake")
+
+    files = []
+    for fname in sorted(os.listdir(real_dir)):
+        if fname.endswith(".wav"):
+            files.append((os.path.join(real_dir, fname), 0, fname))
+    for fname in sorted(os.listdir(fake_dir)):
+        if fname.endswith(".wav"):
+            files.append((os.path.join(fake_dir, fname), 1, fname))
+
+    n_real = sum(1 for _, label, _ in files if label == 0)
+    n_fake = sum(1 for _, label, _ in files if label == 1)
+    print(f"Total files: {len(files)} (real: {n_real}, fake: {n_fake})")
+    print()
+
+    # Run inference
+    results = []
+    total_time = 0.0
+
+    print(
+        f"{'File':<20} {'True':>5} {'Pred':>5} {'P(fake)':>8} "
+        f"{'P(real)':>8} {'Time':>6}"
+    )
+    print("-" * 60)
+
+    for path, true_label, fname in files:
+        audio, sr = librosa.load(path, sr=SAMPLE_RATE, mono=True)
+        audio = pad_audio(audio)
+        x = torch.FloatTensor(audio).unsqueeze(0)
+
+        t0 = time.time()
+        with torch.no_grad():
+            out = ssl_model(x)
+        elapsed = time.time() - t0
+        total_time += elapsed
+
+        probs = torch.softmax(out, dim=1)
+        p_fake = probs[0, 0].item()
+        p_real = probs[0, 1].item()
+        pred = 1 if p_fake >= 0.5 else 0
+
+        results.append(
+            {
+                "file": fname,
+                "true_label": true_label,
+                "pred_label": pred,
+                "prob_fake": p_fake,
+                "prob_real": p_real,
+            }
+        )
+
+        true_str = "FAKE" if true_label == 1 else "REAL"
+        pred_str = "FAKE" if pred == 1 else "REAL"
+        correct = "ok" if pred == true_label else "XX"
+        print(
+            f"{fname:<20} {true_str:>5} {pred_str:>5} "
+            f"{p_fake:>8.4f} {p_real:>8.4f} {elapsed:>5.2f}s "
+            f"[{correct}]"
+        )
+
+    print()
+    print("=" * 70)
+    print("METRICS")
+    print("=" * 70)
+
+    # Compute metrics
+    true_labels = np.array([r["true_label"] for r in results])
+    pred_labels = np.array([r["pred_label"] for r in results])
+
+    tp = int(np.sum((pred_labels == 1) & (true_labels == 1)))
+    tn = int(np.sum((pred_labels == 0) & (true_labels == 0)))
+    fp = int(np.sum((pred_labels == 1) & (true_labels == 0)))
+    fn = int(np.sum((pred_labels == 0) & (true_labels == 1)))
+
+    accuracy = (tp + tn) / len(results) if len(results) > 0 else 0
+    precision = tp / (tp + fp) if (tp + fp) > 0 else 0
+    recall = tp / (tp + fn) if (tp + fn) > 0 else 0
+    f1 = (
+        2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0
+    )
+    specificity = tn / (tn + fp) if (tn + fp) > 0 else 0
+
+    # EER using bonafide scores (prob_real: higher = more real)
+    bonafide_scores = np.array(
+        [r["prob_real"] for r in results if r["true_label"] == 0]
+    )
+    spoof_scores = np.array([r["prob_real"] for r in results if r["true_label"] == 1])
+    if len(bonafide_scores) > 0 and len(spoof_scores) > 0:
+        eer = compute_eer(bonafide_scores, spoof_scores)
+    else:
+        eer = float("nan")
+
+    print(f"Accuracy:    {accuracy:.4f} ({accuracy * 100:.1f}%)")
+    print(f"Precision:   {precision:.4f}")
+    print(f"Recall:      {recall:.4f}")
+    print(f"F1 Score:    {f1:.4f}")
+    print(f"Specificity: {specificity:.4f}")
+    print(f"EER:         {eer:.4f} ({eer * 100:.1f}%)")
+    print()
+    print(f"Confusion Matrix:")
+    print(f"  TP={tp:>3d}  FP={fp:>3d}")
+    print(f"  FN={fn:>3d}  TN={tn:>3d}")
+    print()
+    print(f"Total inference time: {total_time:.1f}s")
+    print(f"Avg per file: {total_time / len(results):.3f}s")
+    print(f"Total files: {len(results)}")
+
+
+if __name__ == "__main__":
+    main()

audio/shiftyspeech/requirements.txt

@@ -0,0 +1,14 @@
+fastapi
+uvicorn
+python-multipart
+torch==2.5.1
+torchaudio==2.5.1
+numpy==1.23.5
+scipy
+librosa==0.9.1
+soundfile
+pydantic
+omegaconf
+hydra-core
+scikit-learn
+bitarray

audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/.env

@@ -0,0 +1,2 @@
+WANDB_API_KEY="<wandb api-key>"
+WANDB_PROJECT_NAME="SSL-AASIST"

audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 Hemlata
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/RawBoost.py

@@ -0,0 +1,143 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+import copy
+
+import numpy as np
+from scipy import signal
+
+"""
+   Hemlata Tak, Madhu Kamble, Jose Patino, Massimiliano Todisco, Nicholas Evans.
+   RawBoost: A Raw Data Boosting and Augmentation Method applied to Automatic Speaker Verification Anti-Spoofing.
+   In Proc. ICASSP 2022, pp:6382--6386.
+"""
+
+
+def randRange(x1, x2, integer):
+    y = np.random.uniform(low=x1, high=x2, size=(1,))
+    if integer:
+        y = int(y)
+    return y
+
+
+def normWav(x, always):
+    if always:
+        x = x / np.amax(abs(x))
+    elif np.amax(abs(x)) > 1:
+        x = x / np.amax(abs(x))
+    return x
+
+
+def genNotchCoeffs(
+    nBands, minF, maxF, minBW, maxBW, minCoeff, maxCoeff, minG, maxG, fs
+):
+    b = 1
+    for i in range(0, nBands):
+        fc = randRange(minF, maxF, 0)
+        bw = randRange(minBW, maxBW, 0)
+        c = randRange(minCoeff, maxCoeff, 1)
+
+        if c / 2 == int(c / 2):
+            c = c + 1
+        f1 = fc - bw / 2
+        f2 = fc + bw / 2
+        if f1 <= 0:
+            f1 = 1 / 1000
+        if f2 >= fs / 2:
+            f2 = fs / 2 - 1 / 1000
+        b = np.convolve(
+            signal.firwin(c, [float(f1), float(f2)], window="hamming", fs=fs), b
+        )
+
+    G = randRange(minG, maxG, 0)
+    _, h = signal.freqz(b, 1, fs=fs)
+    b = pow(10, G / 20) * b / np.amax(abs(h))
+    return b
+
+
+def filterFIR(x, b):
+    N = b.shape[0] + 1
+    xpad = np.pad(x, (0, N), "constant")
+    y = signal.lfilter(b, 1, xpad)
+    y = y[int(N / 2) : int(y.shape[0] - N / 2)]
+    return y
+
+
+# Linear and non-linear convolutive noise
+def LnL_convolutive_noise(
+    x,
+    N_f,
+    nBands,
+    minF,
+    maxF,
+    minBW,
+    maxBW,
+    minCoeff,
+    maxCoeff,
+    minG,
+    maxG,
+    minBiasLinNonLin,
+    maxBiasLinNonLin,
+    fs,
+):
+    y = [0] * x.shape[0]
+    for i in range(0, N_f):
+        if i == 1:
+            minG = minG - minBiasLinNonLin
+            maxG = maxG - maxBiasLinNonLin
+        b = genNotchCoeffs(
+            nBands, minF, maxF, minBW, maxBW, minCoeff, maxCoeff, minG, maxG, fs
+        )
+        y = y + filterFIR(np.power(x, (i + 1)), b)
+    y = y - np.mean(y)
+    y = normWav(y, 0)
+    return y
+
+
+# Impulsive signal dependent noise
+def ISD_additive_noise(x, P, g_sd):
+    beta = randRange(0, P, 0)
+
+    y = copy.deepcopy(x)
+    x_len = x.shape[0]
+    n = int(x_len * (beta / 100))
+    p = np.random.permutation(x_len)[:n]
+    f_r = np.multiply(
+        ((2 * np.random.rand(p.shape[0])) - 1), ((2 * np.random.rand(p.shape[0])) - 1)
+    )
+    r = g_sd * x[p] * f_r
+    y[p] = x[p] + r
+    y = normWav(y, 0)
+    return y
+
+
+# Stationary signal independent noise
+
+
+def SSI_additive_noise(
+    x,
+    SNRmin,
+    SNRmax,
+    nBands,
+    minF,
+    maxF,
+    minBW,
+    maxBW,
+    minCoeff,
+    maxCoeff,
+    minG,
+    maxG,
+    fs,
+):
+    noise = np.random.normal(0, 1, x.shape[0])
+    b = genNotchCoeffs(
+        nBands, minF, maxF, minBW, maxBW, minCoeff, maxCoeff, minG, maxG, fs
+    )
+    noise = filterFIR(noise, b)
+    noise = normWav(noise, 1)
+    SNR = randRange(SNRmin, SNRmax, 0)
+    noise = (
+        noise / np.linalg.norm(noise, 2) * np.linalg.norm(x, 2) / 10.0 ** (0.05 * SNR)
+    )
+    x = x + noise
+    return x

audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/Simplified_CM_solution.py

@@ -0,0 +1,227 @@
+import math
+from collections import OrderedDict
+
+import fairseq
+import numpy as np
+import scipy.io as sio
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from torch.autograd import Variable
+from torch.nn.parameter import Parameter
+from torch.utils import data
+
+___author__ = "Hemlata Tak"
+__email__ = "tak@eurecom.fr"
+
+# from losses_anti_spoofing import AMSoftmax
+
+############################
+## FOR fine-tuning SSL MODEL
+############################
+
+
+class SSLModel(nn.Module):
+    def __init__(self, device):
+        super(SSLModel, self).__init__()
+
+        cp_path = "/change_to_path_to_pre_trained_model_XLR_300M/xlsr2_300m.pt"
+        model, cfg, task = fairseq.checkpoint_utils.load_model_ensemble_and_task(
+            [cp_path]
+        )
+        self.model = model[0]
+        self.device = device
+        self.out_dim = 1024
+        return
+
+    def extract_feat(self, input_data):
+
+        # put the model to GPU if it not there
+        if (
+            next(self.model.parameters()).device != input_data.device
+            or next(self.model.parameters()).dtype != input_data.dtype
+        ):
+            self.model.to(input_data.device, dtype=input_data.dtype)
+            self.model.train()
+
+        if True:
+            # input should be in shape (batch, length)
+            if input_data.ndim == 3:
+                input_tmp = input_data[:, :, 0]
+            else:
+                input_tmp = input_data
+
+            # [batch, length, dim]
+            emb = self.model(input_tmp, mask=False, features_only=True)["x"]
+        return emb
+
+
+# ---------Graph attention simple back-end------------------------#
+""" 
+    Hemlata Tak, Jee-weon Jung, Jose Patino, Madhu Kamble, Massimiliano Todisco, Nicholas Evans.
+    End-to-end spectro-temporal graph attention networks for speaker verification anti-spoofing and speech deepfake detection.
+    In Proc. Automatic Speaker Verification and Spoofing Countermeasures Challenge 2021 Interspeech 2021 satellite workshop.
+"""
+
+
+class GraphAttentionLayer(nn.Module):
+    def __init__(self, in_dim, out_dim, **kwargs):
+        super(GraphAttentionLayer, self).__init__()
+
+        # attention map
+        self.att_proj = nn.Linear(in_dim, out_dim)
+        self.att_weight = self._init_new_params(out_dim, 1)
+
+        # project
+        self.proj_with_att = nn.Linear(in_dim, out_dim)
+        self.proj_without_att = nn.Linear(in_dim, out_dim)
+
+        # batch norm
+        self.bn = nn.BatchNorm1d(out_dim)
+
+        # dropout for inputs
+        self.input_drop = nn.Dropout(p=0.2)
+
+        self.act = nn.SELU(inplace=True)
+
+    def forward(self, x):
+        """
+        x   :(#bs, #node, #dim)
+        """
+        # apply input dropout
+        x = self.input_drop(x)
+
+        # derive attention map
+        att_map = self._derive_att_map(x)
+
+        # projection
+        x = self._project(x, att_map)
+
+        # apply batch norm
+        x = self._apply_BN(x)
+        x = self.act(x)
+
+        return x
+
+    def _pairwise_mul_nodes(self, x):
+        """
+        Calculates pairwise multiplication of nodes.
+        - for attention map
+        x           :(#bs, #node, #dim)
+        out_shape   :(#bs, #node, #node, #dim)
+        """
+
+        nb_nodes = x.size(1)
+        x = x.unsqueeze(2).expand(-1, -1, nb_nodes, -1)
+        x_mirror = x.transpose(1, 2)
+
+        return x * x_mirror
+
+    def _derive_att_map(self, x):
+        """
+        x           :(#bs, #node, #dim)
+        out_shape   :(#bs, #node, #node, 1)
+        """
+        att_map = self._pairwise_mul_nodes(x)
+        att_map = torch.tanh(
+            self.att_proj(att_map)
+        )  # size: (#bs, #node, #node, #dim_out)
+        att_map = torch.matmul(att_map, self.att_weight)  # size: (#bs, #node, #node, 1)
+        att_map = F.softmax(att_map, dim=-2)
+
+        return att_map
+
+    def _project(self, x, att_map):
+        x1 = self.proj_with_att(torch.matmul(att_map.squeeze(-1), x))
+        x2 = self.proj_without_att(x)
+
+        return x1 + x2
+
+    def _apply_BN(self, x):
+        org_size = x.size()
+        x = x.view(-1, org_size[-1])
+        x = self.bn(x)
+        x = x.view(org_size)
+
+        return x
+
+    def _init_new_params(self, *size):
+        out = nn.Parameter(torch.FloatTensor(*size))
+        nn.init.xavier_normal_(out)
+        return out
+
+
+class GraphPool(nn.Module):
+    def __init__(self, k: float, in_dim: int, p):
+        super().__init__()
+        self.k = k
+        self.sigmoid = nn.Sigmoid()
+        self.proj = nn.Linear(in_dim, 1)
+        self.drop = nn.Dropout(p=p) if p > 0 else nn.Identity()
+        self.in_dim = in_dim
+
+    def forward(self, h):
+        Z = self.drop(h)
+        weights = self.proj(Z)
+        scores = self.sigmoid(weights)
+        new_h = self.top_k_graph(scores, h, self.k)
+
+        return new_h
+
+    def top_k_graph(self, scores, h, k):
+        """
+        args
+        =====
+        scores: attention-based weights (#bs, #node, 1)
+        h: graph data (#bs, #node, #dim)
+        k: ratio of remaining nodes, (float)
+        returns
+        =====
+        h: graph pool applied data (#bs, #node', #dim)
+        """
+        _, n_nodes, n_feat = h.size()
+        n_nodes = max(int(n_nodes * k), 1)
+        _, idx = torch.topk(scores, n_nodes, dim=1)
+        idx = idx.expand(-1, -1, n_feat)
+
+        h = h * scores
+        h = torch.gather(h, 1, idx)
+
+        return h
+
+
+class Model(nn.Module):
+    def __init__(self, d_args, device):
+        super(Model, self).__init__()
+
+        # SSL model
+        self.device = device
+        self.ssl_model = SSLModel(self.device)
+        self.LL = nn.Linear(self.ssl_model.out_dim, 128)
+        self.first_bn = nn.BatchNorm1d(num_features=128)
+        self.selu = nn.SELU(inplace=True)
+
+        # graph module layer
+        self.GAT_layer = GraphAttentionLayer(128, 64)
+        self.proj = nn.Linear(64, 1)
+        self.pool = GraphPool(0.8, 64, 0.3)
+
+        # classifier head
+        self.proj_node = nn.Linear(53, 2)
+
+    def forward(self, x_inp, Freq_aug=False):
+        # SSL wav2vec 2.0 model
+        x_ssl_feat = self.ssl_model.extract_feat(x_inp.squeeze(-1))
+        x_SSL = self.LL(x_ssl_feat)  # (bs,frame_number,feat_out_dim)
+        x_SSL = x_SSL.transpose(1, 2)  # (bs,feat_out_dim,frame_number)
+
+        x = F.max_pool1d(x_SSL, (3))
+        x = self.first_bn(x)
+        x = self.selu(x)
+
+        x = self.GAT_layer(x.transpose(1, 2))
+        x = self.pool(x)
+        x = self.proj(x).flatten(1)
+        output = self.proj_node(x)
+        return output

audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/data_utils.py

@@ -0,0 +1,292 @@
+import os
+import random
+from random import randrange
+
+import librosa
+import numpy as np
+import torch
+import torch.nn as nn
+from RawBoost import (
+    ISD_additive_noise,
+    LnL_convolutive_noise,
+    SSI_additive_noise,
+    normWav,
+)
+from torch import Tensor
+from torch.utils.data import Dataset
+
+__author__ = "Hemlata Tak"
+__email__ = "tak@eurecom.fr"
+
+
+def genSpoof_list(dir_meta, is_train=False, is_eval=False):
+    d_meta = {}
+    file_list = []
+    with open(dir_meta, "r") as f:
+        l_meta = f.readlines()
+
+    if is_train:
+        for line in l_meta:
+            key, label = line.strip().split()
+            file_list.append(key)
+            d_meta[key] = 1 if label == "bonafide" else 0
+        return d_meta, file_list
+
+    elif is_eval:
+        for line in l_meta:
+            key, _ = line.strip().split(" ")
+            file_list.append(key)
+        return file_list
+    else:
+        for line in l_meta:
+            key, label = line.strip().split()
+            file_list.append(key)
+            d_meta[key] = 1 if label == "bonafide" else 0
+        return d_meta, file_list
+
+
+def pad(x, max_len=64600):
+    x_len = x.shape[0]
+    if x_len >= max_len:
+        return x[:max_len]
+    # need to pad
+    num_repeats = int(max_len / x_len) + 1
+    padded_x = np.tile(x, (1, num_repeats))[:, :max_len][0]
+    return padded_x
+
+
+class Dataset_ASVspoof2019_train(Dataset):
+    def __init__(self, args, metafile, algo):
+        """self.list_IDs : list of strings (each string: utt key),
+        self.labels: dictionary (key: utt key, value: label integer)"""
+
+        self.uttpath_labels = []
+        with open(metafile, "r") as f:
+            for line in f:
+                items = line.strip().split()
+                lb = 1 if items[-1] == "bonafide" else 0
+                self.uttpath_labels.append((items[0], lb))
+
+        self.algo = algo
+        self.args = args
+        self.cut = 64600  # take ~4 sec audio (64600 samples)
+
+    def __len__(self):
+        return len(self.uttpath_labels)
+
+    def __getitem__(self, index):
+        path, target = self.uttpath_labels[index]
+        X, fs = librosa.load(path, sr=16000)
+        Y = process_Rawboost_feature(X, fs, self.args, self.algo)
+        X_pad = pad(Y, self.cut)
+        x_inp = Tensor(X_pad)
+        return x_inp, target
+
+
+class Dataset_ASVspoof2021_eval(Dataset):
+    def __init__(self, list_IDs):
+        """self.list_IDs	: list of strings (each string: utt key),"""
+
+        self.list_IDs = list_IDs
+        self.cut = 64600  # take ~4 sec audio (64600 samples)
+
+    def __len__(self):
+        return len(self.list_IDs)
+
+    def __getitem__(self, index):
+        utt_id = self.list_IDs[index]
+        X, fs = librosa.load(utt_id, sr=16000)
+        X_pad = pad(X, self.cut)
+        x_inp = Tensor(X_pad)
+        return x_inp, utt_id
+
+
+# --------------RawBoost data augmentation algorithms---------------------------##
+def process_Rawboost_feature(feature, sr, args, algo):
+
+    # Data process by Convolutive noise (1st algo)
+    if algo == 1:
+
+        feature = LnL_convolutive_noise(
+            feature,
+            args.N_f,
+            args.nBands,
+            args.minF,
+            args.maxF,
+            args.minBW,
+            args.maxBW,
+            args.minCoeff,
+            args.maxCoeff,
+            args.minG,
+            args.maxG,
+            args.minBiasLinNonLin,
+            args.maxBiasLinNonLin,
+            sr,
+        )
+
+    # Data process by Impulsive noise (2nd algo)
+    elif algo == 2:
+
+        feature = ISD_additive_noise(feature, args.P, args.g_sd)
+
+    # Data process by coloured additive noise (3rd algo)
+    elif algo == 3:
+
+        feature = SSI_additive_noise(
+            feature,
+            args.SNRmin,
+            args.SNRmax,
+            args.nBands,
+            args.minF,
+            args.maxF,
+            args.minBW,
+            args.maxBW,
+            args.minCoeff,
+            args.maxCoeff,
+            args.minG,
+            args.maxG,
+            sr,
+        )
+
+    # Data process by all 3 algo. together in series (1+2+3)
+    elif algo == 4:
+
+        feature = LnL_convolutive_noise(
+            feature,
+            args.N_f,
+            args.nBands,
+            args.minF,
+            args.maxF,
+            args.minBW,
+            args.maxBW,
+            args.minCoeff,
+            args.maxCoeff,
+            args.minG,
+            args.maxG,
+            args.minBiasLinNonLin,
+            args.maxBiasLinNonLin,
+            sr,
+        )
+        feature = ISD_additive_noise(feature, args.P, args.g_sd)
+        feature = SSI_additive_noise(
+            feature,
+            args.SNRmin,
+            args.SNRmax,
+            args.nBands,
+            args.minF,
+            args.maxF,
+            args.minBW,
+            args.maxBW,
+            args.minCoeff,
+            args.maxCoeff,
+            args.minG,
+            args.maxG,
+            sr,
+        )
+
+    # Data process by 1st two algo. together in series (1+2)
+    elif algo == 5:
+
+        feature = LnL_convolutive_noise(
+            feature,
+            args.N_f,
+            args.nBands,
+            args.minF,
+            args.maxF,
+            args.minBW,
+            args.maxBW,
+            args.minCoeff,
+            args.maxCoeff,
+            args.minG,
+            args.maxG,
+            args.minBiasLinNonLin,
+            args.maxBiasLinNonLin,
+            sr,
+        )
+        feature = ISD_additive_noise(feature, args.P, args.g_sd)
+
+    # Data process by 1st and 3rd algo. together in series (1+3)
+    elif algo == 6:
+
+        feature = LnL_convolutive_noise(
+            feature,
+            args.N_f,
+            args.nBands,
+            args.minF,
+            args.maxF,
+            args.minBW,
+            args.maxBW,
+            args.minCoeff,
+            args.maxCoeff,
+            args.minG,
+            args.maxG,
+            args.minBiasLinNonLin,
+            args.maxBiasLinNonLin,
+            sr,
+        )
+        feature = SSI_additive_noise(
+            feature,
+            args.SNRmin,
+            args.SNRmax,
+            args.nBands,
+            args.minF,
+            args.maxF,
+            args.minBW,
+            args.maxBW,
+            args.minCoeff,
+            args.maxCoeff,
+            args.minG,
+            args.maxG,
+            sr,
+        )
+
+    # Data process by 2nd and 3rd algo. together in series (2+3)
+    elif algo == 7:
+
+        feature = ISD_additive_noise(feature, args.P, args.g_sd)
+        feature = SSI_additive_noise(
+            feature,
+            args.SNRmin,
+            args.SNRmax,
+            args.nBands,
+            args.minF,
+            args.maxF,
+            args.minBW,
+            args.maxBW,
+            args.minCoeff,
+            args.maxCoeff,
+            args.minG,
+            args.maxG,
+            sr,
+        )
+
+    # Data process by 1st two algo. together in Parallel (1||2)
+    elif algo == 8:
+
+        feature1 = LnL_convolutive_noise(
+            feature,
+            args.N_f,
+            args.nBands,
+            args.minF,
+            args.maxF,
+            args.minBW,
+            args.maxBW,
+            args.minCoeff,
+            args.maxCoeff,
+            args.minG,
+            args.maxG,
+            args.minBiasLinNonLin,
+            args.maxBiasLinNonLin,
+            sr,
+        )
+        feature2 = ISD_additive_noise(feature, args.P, args.g_sd)
+
+        feature_para = feature1 + feature2
+        feature = normWav(feature_para, 0)  # normalized resultant waveform
+
+    # original data without Rawboost processing
+    else:
+
+        feature = feature
+
+    return feature

audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/model.py

@@ -0,0 +1,603 @@
+import random
+from typing import Union
+
+import fairseq
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+___author__ = "Hemlata Tak"
+__email__ = "tak@eurecom.fr"
+
+############################
+## FOR fine-tuned SSL MODEL
+############################
+
+
+class SSLModel(nn.Module):
+    def __init__(self, device):
+        super(SSLModel, self).__init__()
+
+        cp_path = "models/xlsr2_300m.pt"  # Change the pre-trained XLSR model path.
+        model, cfg, task = fairseq.checkpoint_utils.load_model_ensemble_and_task(
+            [cp_path]
+        )
+        self.model = model[0]
+        self.device = device
+        self.out_dim = 1024
+        return
+
+    def extract_feat(self, input_data):
+
+        # put the model to GPU if it not there
+        if (
+            next(self.model.parameters()).device != input_data.device
+            or next(self.model.parameters()).dtype != input_data.dtype
+        ):
+            self.model.to(input_data.device, dtype=input_data.dtype)
+            self.model.train()
+
+        if True:
+            # input should be in shape (batch, length)
+            if input_data.ndim == 3:
+                input_tmp = input_data[:, :, 0]
+            else:
+                input_tmp = input_data
+
+            # [batch, length, dim]
+            emb = self.model(input_tmp, mask=False, features_only=True)["x"]
+        return emb
+
+
+# ---------AASIST back-end------------------------#
+""" Jee-weon Jung, Hee-Soo Heo, Hemlata Tak, Hye-jin Shim, Joon Son Chung, Bong-Jin Lee, Ha-Jin Yu and Nicholas Evans. 
+    AASIST: Audio Anti-Spoofing Using Integrated Spectro-Temporal Graph Attention Networks. 
+    In Proc. ICASSP 2022, pp: 6367--6371."""
+
+
+class GraphAttentionLayer(nn.Module):
+    def __init__(self, in_dim, out_dim, **kwargs):
+        super().__init__()
+
+        # attention map
+        self.att_proj = nn.Linear(in_dim, out_dim)
+        self.att_weight = self._init_new_params(out_dim, 1)
+
+        # project
+        self.proj_with_att = nn.Linear(in_dim, out_dim)
+        self.proj_without_att = nn.Linear(in_dim, out_dim)
+
+        # batch norm
+        self.bn = nn.BatchNorm1d(out_dim)
+
+        # dropout for inputs
+        self.input_drop = nn.Dropout(p=0.2)
+
+        # activate
+        self.act = nn.SELU(inplace=True)
+
+        # temperature
+        self.temp = 1.0
+        if "temperature" in kwargs:
+            self.temp = kwargs["temperature"]
+
+    def forward(self, x):
+        """
+        x   :(#bs, #node, #dim)
+        """
+        # apply input dropout
+        x = self.input_drop(x)
+
+        # derive attention map
+        att_map = self._derive_att_map(x)
+
+        # projection
+        x = self._project(x, att_map)
+
+        # apply batch norm
+        x = self._apply_BN(x)
+        x = self.act(x)
+        return x
+
+    def _pairwise_mul_nodes(self, x):
+        """
+        Calculates pairwise multiplication of nodes.
+        - for attention map
+        x           :(#bs, #node, #dim)
+        out_shape   :(#bs, #node, #node, #dim)
+        """
+
+        nb_nodes = x.size(1)
+        x = x.unsqueeze(2).expand(-1, -1, nb_nodes, -1)
+        x_mirror = x.transpose(1, 2)
+
+        return x * x_mirror
+
+    def _derive_att_map(self, x):
+        """
+        x           :(#bs, #node, #dim)
+        out_shape   :(#bs, #node, #node, 1)
+        """
+        att_map = self._pairwise_mul_nodes(x)
+        # size: (#bs, #node, #node, #dim_out)
+        att_map = torch.tanh(self.att_proj(att_map))
+        # size: (#bs, #node, #node, 1)
+        att_map = torch.matmul(att_map, self.att_weight)
+
+        # apply temperature
+        att_map = att_map / self.temp
+
+        att_map = F.softmax(att_map, dim=-2)
+
+        return att_map
+
+    def _project(self, x, att_map):
+        x1 = self.proj_with_att(torch.matmul(att_map.squeeze(-1), x))
+        x2 = self.proj_without_att(x)
+
+        return x1 + x2
+
+    def _apply_BN(self, x):
+        org_size = x.size()
+        x = x.view(-1, org_size[-1])
+        x = self.bn(x)
+        x = x.view(org_size)
+
+        return x
+
+    def _init_new_params(self, *size):
+        out = nn.Parameter(torch.FloatTensor(*size))
+        nn.init.xavier_normal_(out)
+        return out
+
+
+class HtrgGraphAttentionLayer(nn.Module):
+    def __init__(self, in_dim, out_dim, **kwargs):
+        super().__init__()
+
+        self.proj_type1 = nn.Linear(in_dim, in_dim)
+        self.proj_type2 = nn.Linear(in_dim, in_dim)
+
+        # attention map
+        self.att_proj = nn.Linear(in_dim, out_dim)
+        self.att_projM = nn.Linear(in_dim, out_dim)
+
+        self.att_weight11 = self._init_new_params(out_dim, 1)
+        self.att_weight22 = self._init_new_params(out_dim, 1)
+        self.att_weight12 = self._init_new_params(out_dim, 1)
+        self.att_weightM = self._init_new_params(out_dim, 1)
+
+        # project
+        self.proj_with_att = nn.Linear(in_dim, out_dim)
+        self.proj_without_att = nn.Linear(in_dim, out_dim)
+
+        self.proj_with_attM = nn.Linear(in_dim, out_dim)
+        self.proj_without_attM = nn.Linear(in_dim, out_dim)
+
+        # batch norm
+        self.bn = nn.BatchNorm1d(out_dim)
+
+        # dropout for inputs
+        self.input_drop = nn.Dropout(p=0.2)
+
+        # activate
+        self.act = nn.SELU(inplace=True)
+
+        # temperature
+        self.temp = 1.0
+        if "temperature" in kwargs:
+            self.temp = kwargs["temperature"]
+
+    def forward(self, x1, x2, master=None):
+        """
+        x1  :(#bs, #node, #dim)
+        x2  :(#bs, #node, #dim)
+        """
+
+        num_type1 = x1.size(1)
+        num_type2 = x2.size(1)
+
+        x1 = self.proj_type1(x1)
+
+        x2 = self.proj_type2(x2)
+
+        x = torch.cat([x1, x2], dim=1)
+
+        if master is None:
+            master = torch.mean(x, dim=1, keepdim=True)
+
+        # apply input dropout
+        x = self.input_drop(x)
+
+        # derive attention map
+        att_map = self._derive_att_map(x, num_type1, num_type2)
+
+        # directional edge for master node
+        master = self._update_master(x, master)
+
+        # projection
+        x = self._project(x, att_map)
+
+        # apply batch norm
+        x = self._apply_BN(x)
+        x = self.act(x)
+
+        x1 = x.narrow(1, 0, num_type1)
+
+        x2 = x.narrow(1, num_type1, num_type2)
+
+        return x1, x2, master
+
+    def _update_master(self, x, master):
+
+        att_map = self._derive_att_map_master(x, master)
+        master = self._project_master(x, master, att_map)
+
+        return master
+
+    def _pairwise_mul_nodes(self, x):
+        """
+        Calculates pairwise multiplication of nodes.
+        - for attention map
+        x           :(#bs, #node, #dim)
+        out_shape   :(#bs, #node, #node, #dim)
+        """
+
+        nb_nodes = x.size(1)
+        x = x.unsqueeze(2).expand(-1, -1, nb_nodes, -1)
+        x_mirror = x.transpose(1, 2)
+
+        return x * x_mirror
+
+    def _derive_att_map_master(self, x, master):
+        """
+        x           :(#bs, #node, #dim)
+        out_shape   :(#bs, #node, #node, 1)
+        """
+        att_map = x * master
+        att_map = torch.tanh(self.att_projM(att_map))
+
+        att_map = torch.matmul(att_map, self.att_weightM)
+
+        # apply temperature
+        att_map = att_map / self.temp
+
+        att_map = F.softmax(att_map, dim=-2)
+
+        return att_map
+
+    def _derive_att_map(self, x, num_type1, num_type2):
+        """
+        x           :(#bs, #node, #dim)
+        out_shape   :(#bs, #node, #node, 1)
+        """
+        att_map = self._pairwise_mul_nodes(x)
+        # size: (#bs, #node, #node, #dim_out)
+        att_map = torch.tanh(self.att_proj(att_map))
+        # size: (#bs, #node, #node, 1)
+
+        att_board = torch.zeros_like(att_map[:, :, :, 0]).unsqueeze(-1)
+
+        att_board[:, :num_type1, :num_type1, :] = torch.matmul(
+            att_map[:, :num_type1, :num_type1, :], self.att_weight11
+        )
+        att_board[:, num_type1:, num_type1:, :] = torch.matmul(
+            att_map[:, num_type1:, num_type1:, :], self.att_weight22
+        )
+        att_board[:, :num_type1, num_type1:, :] = torch.matmul(
+            att_map[:, :num_type1, num_type1:, :], self.att_weight12
+        )
+        att_board[:, num_type1:, :num_type1, :] = torch.matmul(
+            att_map[:, num_type1:, :num_type1, :], self.att_weight12
+        )
+
+        att_map = att_board
+
+        # apply temperature
+        att_map = att_map / self.temp
+
+        att_map = F.softmax(att_map, dim=-2)
+
+        return att_map
+
+    def _project(self, x, att_map):
+        x1 = self.proj_with_att(torch.matmul(att_map.squeeze(-1), x))
+        x2 = self.proj_without_att(x)
+
+        return x1 + x2
+
+    def _project_master(self, x, master, att_map):
+
+        x1 = self.proj_with_attM(torch.matmul(att_map.squeeze(-1).unsqueeze(1), x))
+        x2 = self.proj_without_attM(master)
+
+        return x1 + x2
+
+    def _apply_BN(self, x):
+        org_size = x.size()
+        x = x.view(-1, org_size[-1])
+        x = self.bn(x)
+        x = x.view(org_size)
+
+        return x
+
+    def _init_new_params(self, *size):
+        out = nn.Parameter(torch.FloatTensor(*size))
+        nn.init.xavier_normal_(out)
+        return out
+
+
+class GraphPool(nn.Module):
+    def __init__(self, k: float, in_dim: int, p: Union[float, int]):
+        super().__init__()
+        self.k = k
+        self.sigmoid = nn.Sigmoid()
+        self.proj = nn.Linear(in_dim, 1)
+        self.drop = nn.Dropout(p=p) if p > 0 else nn.Identity()
+        self.in_dim = in_dim
+
+    def forward(self, h):
+        Z = self.drop(h)
+        weights = self.proj(Z)
+        scores = self.sigmoid(weights)
+        new_h = self.top_k_graph(scores, h, self.k)
+
+        return new_h
+
+    def top_k_graph(self, scores, h, k):
+        """
+        args
+        =====
+        scores: attention-based weights (#bs, #node, 1)
+        h: graph data (#bs, #node, #dim)
+        k: ratio of remaining nodes, (float)
+        returns
+        =====
+        h: graph pool applied data (#bs, #node', #dim)
+        """
+        _, n_nodes, n_feat = h.size()
+        n_nodes = max(int(n_nodes * k), 1)
+        _, idx = torch.topk(scores, n_nodes, dim=1)
+        idx = idx.expand(-1, -1, n_feat)
+
+        h = h * scores
+        h = torch.gather(h, 1, idx)
+
+        return h
+
+
+class Residual_block(nn.Module):
+    def __init__(self, nb_filts, first=False):
+        super().__init__()
+        self.first = first
+
+        if not self.first:
+            self.bn1 = nn.BatchNorm2d(num_features=nb_filts[0])
+        self.conv1 = nn.Conv2d(
+            in_channels=nb_filts[0],
+            out_channels=nb_filts[1],
+            kernel_size=(2, 3),
+            padding=(1, 1),
+            stride=1,
+        )
+        self.selu = nn.SELU(inplace=True)
+
+        self.bn2 = nn.BatchNorm2d(num_features=nb_filts[1])
+        self.conv2 = nn.Conv2d(
+            in_channels=nb_filts[1],
+            out_channels=nb_filts[1],
+            kernel_size=(2, 3),
+            padding=(0, 1),
+            stride=1,
+        )
+
+        if nb_filts[0] != nb_filts[1]:
+            self.downsample = True
+            self.conv_downsample = nn.Conv2d(
+                in_channels=nb_filts[0],
+                out_channels=nb_filts[1],
+                padding=(0, 1),
+                kernel_size=(1, 3),
+                stride=1,
+            )
+
+        else:
+            self.downsample = False
+
+    def forward(self, x):
+        identity = x
+        if not self.first:
+            out = self.bn1(x)
+            out = self.selu(out)
+        else:
+            out = x
+
+        out = self.conv1(x)
+
+        out = self.bn2(out)
+        out = self.selu(out)
+
+        out = self.conv2(out)
+
+        if self.downsample:
+            identity = self.conv_downsample(identity)
+
+        out += identity
+
+        return out
+
+
+class Model(nn.Module):
+    def __init__(self, args, device):
+        super().__init__()
+        self.device = device
+
+        # AASIST parameters
+        filts = [128, [1, 32], [32, 32], [32, 64], [64, 64]]
+        gat_dims = [64, 32]
+        pool_ratios = [0.5, 0.5, 0.5, 0.5]
+        temperatures = [2.0, 2.0, 100.0, 100.0]
+
+        ####
+        # create network wav2vec 2.0
+        ####
+        self.ssl_model = SSLModel(self.device)
+        self.LL = nn.Linear(self.ssl_model.out_dim, 128)
+
+        self.first_bn = nn.BatchNorm2d(num_features=1)
+        self.first_bn1 = nn.BatchNorm2d(num_features=64)
+        self.drop = nn.Dropout(0.5, inplace=True)
+        self.drop_way = nn.Dropout(0.2, inplace=True)
+        self.selu = nn.SELU(inplace=True)
+
+        # RawNet2 encoder
+        self.encoder = nn.Sequential(
+            nn.Sequential(Residual_block(nb_filts=filts[1], first=True)),
+            nn.Sequential(Residual_block(nb_filts=filts[2])),
+            nn.Sequential(Residual_block(nb_filts=filts[3])),
+            nn.Sequential(Residual_block(nb_filts=filts[4])),
+            nn.Sequential(Residual_block(nb_filts=filts[4])),
+            nn.Sequential(Residual_block(nb_filts=filts[4])),
+        )
+
+        self.attention = nn.Sequential(
+            nn.Conv2d(64, 128, kernel_size=(1, 1)),
+            nn.SELU(inplace=True),
+            nn.BatchNorm2d(128),
+            nn.Conv2d(128, 64, kernel_size=(1, 1)),
+        )
+        # position encoding
+        self.pos_S = nn.Parameter(torch.randn(1, 42, filts[-1][-1]))
+
+        self.master1 = nn.Parameter(torch.randn(1, 1, gat_dims[0]))
+        self.master2 = nn.Parameter(torch.randn(1, 1, gat_dims[0]))
+
+        # Graph module
+        self.GAT_layer_S = GraphAttentionLayer(
+            filts[-1][-1], gat_dims[0], temperature=temperatures[0]
+        )
+        self.GAT_layer_T = GraphAttentionLayer(
+            filts[-1][-1], gat_dims[0], temperature=temperatures[1]
+        )
+        # HS-GAL layer
+        self.HtrgGAT_layer_ST11 = HtrgGraphAttentionLayer(
+            gat_dims[0], gat_dims[1], temperature=temperatures[2]
+        )
+        self.HtrgGAT_layer_ST12 = HtrgGraphAttentionLayer(
+            gat_dims[1], gat_dims[1], temperature=temperatures[2]
+        )
+        self.HtrgGAT_layer_ST21 = HtrgGraphAttentionLayer(
+            gat_dims[0], gat_dims[1], temperature=temperatures[2]
+        )
+        self.HtrgGAT_layer_ST22 = HtrgGraphAttentionLayer(
+            gat_dims[1], gat_dims[1], temperature=temperatures[2]
+        )
+
+        # Graph pooling layers
+        self.pool_S = GraphPool(pool_ratios[0], gat_dims[0], 0.3)
+        self.pool_T = GraphPool(pool_ratios[1], gat_dims[0], 0.3)
+        self.pool_hS1 = GraphPool(pool_ratios[2], gat_dims[1], 0.3)
+        self.pool_hT1 = GraphPool(pool_ratios[2], gat_dims[1], 0.3)
+
+        self.pool_hS2 = GraphPool(pool_ratios[2], gat_dims[1], 0.3)
+        self.pool_hT2 = GraphPool(pool_ratios[2], gat_dims[1], 0.3)
+
+        self.out_layer = nn.Linear(5 * gat_dims[1], 2)
+
+    def forward(self, x):
+        # -------pre-trained Wav2vec model fine tunning ------------------------##
+        x_ssl_feat = self.ssl_model.extract_feat(x.squeeze(-1))
+        x = self.LL(x_ssl_feat)  # (bs,frame_number,feat_out_dim)
+
+        # post-processing on front-end features
+        x = x.transpose(1, 2)  # (bs,feat_out_dim,frame_number)
+        x = x.unsqueeze(dim=1)  # add channel
+        x = F.max_pool2d(x, (3, 3))
+        x = self.first_bn(x)
+        x = self.selu(x)
+
+        # RawNet2-based encoder
+        x = self.encoder(x)
+        x = self.first_bn1(x)
+        x = self.selu(x)
+
+        w = self.attention(x)
+
+        # ------------SA for spectral feature-------------#
+        w1 = F.softmax(w, dim=-1)
+        m = torch.sum(x * w1, dim=-1)
+        e_S = m.transpose(1, 2) + self.pos_S
+
+        # graph module layer
+        gat_S = self.GAT_layer_S(e_S)
+        out_S = self.pool_S(gat_S)  # (#bs, #node, #dim)
+
+        # ------------SA for temporal feature-------------#
+        w2 = F.softmax(w, dim=-2)
+        m1 = torch.sum(x * w2, dim=-2)
+
+        e_T = m1.transpose(1, 2)
+
+        # graph module layer
+        gat_T = self.GAT_layer_T(e_T)
+        out_T = self.pool_T(gat_T)
+
+        # learnable master node
+        master1 = self.master1.expand(x.size(0), -1, -1)
+        master2 = self.master2.expand(x.size(0), -1, -1)
+
+        # inference 1
+        out_T1, out_S1, master1 = self.HtrgGAT_layer_ST11(
+            out_T, out_S, master=self.master1
+        )
+
+        out_S1 = self.pool_hS1(out_S1)
+        out_T1 = self.pool_hT1(out_T1)
+
+        out_T_aug, out_S_aug, master_aug = self.HtrgGAT_layer_ST12(
+            out_T1, out_S1, master=master1
+        )
+        out_T1 = out_T1 + out_T_aug
+        out_S1 = out_S1 + out_S_aug
+        master1 = master1 + master_aug
+
+        # inference 2
+        out_T2, out_S2, master2 = self.HtrgGAT_layer_ST21(
+            out_T, out_S, master=self.master2
+        )
+        out_S2 = self.pool_hS2(out_S2)
+        out_T2 = self.pool_hT2(out_T2)
+
+        out_T_aug, out_S_aug, master_aug = self.HtrgGAT_layer_ST22(
+            out_T2, out_S2, master=master2
+        )
+        out_T2 = out_T2 + out_T_aug
+        out_S2 = out_S2 + out_S_aug
+        master2 = master2 + master_aug
+
+        out_T1 = self.drop_way(out_T1)
+        out_T2 = self.drop_way(out_T2)
+        out_S1 = self.drop_way(out_S1)
+        out_S2 = self.drop_way(out_S2)
+        master1 = self.drop_way(master1)
+        master2 = self.drop_way(master2)
+
+        out_T = torch.max(out_T1, out_T2)
+        out_S = torch.max(out_S1, out_S2)
+        master = torch.max(master1, master2)
+
+        # Readout operation
+        T_max, _ = torch.max(torch.abs(out_T), dim=1)
+        T_avg = torch.mean(out_T, dim=1)
+
+        S_max, _ = torch.max(torch.abs(out_S), dim=1)
+        S_avg = torch.mean(out_S, dim=1)
+
+        last_hidden = torch.cat([T_max, T_avg, S_max, S_avg, master.squeeze(1)], dim=1)
+
+        last_hidden = self.drop(last_hidden)
+        output = self.out_layer(last_hidden)
+
+        return output

audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/requirements.txt

@@ -0,0 +1,4 @@
+librosa==0.9.1
+python-dotenv==1.0.1
+tensorboardX
+wandb

audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/startup_config.py

@@ -0,0 +1,60 @@
+#!/usr/bin/env python
+"""
+startup_config
+
+Startup configuration utilities
+
+"""
+
+from __future__ import absolute_import
+
+import importlib
+import os
+import random
+import sys
+
+import numpy as np
+import torch
+
+__author__ = "Xin Wang"
+__email__ = "wangxin@nii.ac.jp"
+__copyright__ = "Copyright 2020, Xin Wang"
+
+
+def set_random_seed(random_seed, args=None):
+    """set_random_seed(random_seed, args=None)
+
+    Set the random_seed for numpy, python, and cudnn
+
+    input
+    -----
+      random_seed: integer random seed
+      args: argue parser
+    """
+
+    # initialization
+    torch.manual_seed(random_seed)
+    random.seed(random_seed)
+    np.random.seed(random_seed)
+    os.environ["PYTHONHASHSEED"] = str(random_seed)
+
+    # For torch.backends.cudnn.deterministic
+    # Note: this default configuration may result in RuntimeError
+    # see https://pytorch.org/docs/stable/notes/randomness.html
+    if args is None:
+        cudnn_deterministic = True
+        cudnn_benchmark = False
+    else:
+        cudnn_deterministic = args.cudnn_deterministic_toggle
+        cudnn_benchmark = args.cudnn_benchmark_toggle
+
+        if not cudnn_deterministic:
+            print("cudnn_deterministic set to False")
+        if cudnn_benchmark:
+            print("cudnn_benchmark set to True")
+
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(random_seed)
+        torch.backends.cudnn.deterministic = cudnn_deterministic
+        torch.backends.cudnn.benchmark = cudnn_benchmark
+    return

audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/train.py

@@ -0,0 +1,446 @@
+import argparse
+import os
+import sys
+
+import librosa
+import numpy as np
+import torch
+import wandb
+import yaml
+from data_utils import (
+    Dataset_ASVspoof2019_train,
+    Dataset_ASVspoof2021_eval,
+    genSpoof_list,
+    pad,
+    process_Rawboost_feature,
+)
+from dotenv import load_dotenv
+from model import Model
+from sklearn.metrics import roc_auc_score
+from startup_config import set_random_seed
+from tensorboardX import SummaryWriter
+from torch import Tensor, nn
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+__author__ = "Hemlata Tak"
+__email__ = "tak@eurecom.fr"
+
+
+def compute_det_curve(target_scores, nontarget_scores):
+
+    n_scores = target_scores.size + nontarget_scores.size
+    all_scores = np.concatenate((target_scores, nontarget_scores))
+    labels = np.concatenate(
+        (np.ones(target_scores.size), np.zeros(nontarget_scores.size))
+    )
+
+    indices = np.argsort(all_scores, kind="mergesort")
+    labels = labels[indices]
+    tar_trial_sums = np.cumsum(labels)
+    nontarget_trial_sums = nontarget_scores.size - (
+        np.arange(1, n_scores + 1) - tar_trial_sums
+    )
+
+    frr = np.concatenate((np.atleast_1d(0), tar_trial_sums / target_scores.size))
+    far = np.concatenate(
+        (np.atleast_1d(1), nontarget_trial_sums / nontarget_scores.size)
+    )
+    # Thresholds are the sorted scores
+    thresholds = np.concatenate(
+        (np.atleast_1d(all_scores[indices[0]] - 0.001), all_scores[indices])
+    )
+
+    return frr, far, thresholds
+
+
+def compute_eer(target_scores, nontarget_scores):
+    """Returns equal error rate (EER) and the corresponding threshold."""
+    frr, far, thresholds = compute_det_curve(target_scores, nontarget_scores)
+    abs_diffs = np.abs(frr - far)
+    min_index = np.argmin(abs_diffs)
+    eer = np.mean((frr[min_index], far[min_index]))
+    return eer, thresholds[min_index], frr, far
+
+
+def calculate_tDCF_EER(cm_scores_file, output_file, printout=True):
+    # Load CM scores
+    cm_data = np.genfromtxt(cm_scores_file, dtype=str)
+    cm_utt_id = cm_data[:, 0]
+    cm_keys = cm_data[:, 1]
+    cm_scores = cm_data[:, 2].astype(float)
+    # Extract bona fide (real human) and spoof scores from the CM scores
+    bona_cm = cm_scores[cm_keys == "bonafide"]
+    spoof_cm = cm_scores[cm_keys == "spoof"]
+    all_scores = np.concatenate([bona_cm, spoof_cm])
+    all_true_labels = np.concatenate([np.ones_like(bona_cm), np.zeros_like(spoof_cm)])
+
+    auc = roc_auc_score(all_true_labels, all_scores, max_fpr=0.05)
+    eer_cm, eer_threshold, frr, far = compute_eer(bona_cm, spoof_cm)
+
+    if printout:
+        with open(output_file, "w") as f_res:
+            f_res.write("\nCM SYSTEM\n")
+            f_res.write(
+                "\tEER\t\t= {:8.9f} % "
+                "(Equal error rate for countermeasure)\n".format(eer_cm * 100)
+            )
+            f_res.write("\t pAUC with max fpr - 0.05 is :{}".format(auc))
+
+
+def evaluate_accuracy(dev_loader, model, device, args):
+    val_loss = 0.0
+    num_total = 0.0
+    algo = args.algo
+    cut = 64600
+    model.eval()
+
+    weight = torch.FloatTensor([0.1, 0.9]).to(device)
+    criterion = nn.CrossEntropyLoss(weight=weight)
+    progress_bar = tqdm(dev_loader, desc=f"Epoch {epoch+1}/{args.num_epochs}")
+    for current_step, (batch_pths, batch_y) in enumerate(progress_bar):
+        batch_x = batch_pths
+        batch_size = batch_x.size(0)
+        num_total += batch_size
+        batch_x = batch_x.to(device)
+        batch_y = batch_y.view(-1).type(torch.int64).to(device)
+        batch_out = model(batch_x)
+
+        batch_loss = criterion(batch_out, batch_y)
+        val_loss += batch_loss.item() * batch_size
+
+    val_loss /= num_total
+
+    return val_loss
+
+
+def produce_evaluation_file(dataset, model, device, save_path, trial_path):
+    data_loader = DataLoader(dataset, batch_size=10, shuffle=False, drop_last=False)
+    num_correct = 0.0
+    num_total = 0.0
+    model.eval()
+    with open(trial_path, "r") as f_trl:
+        trial_lines = f_trl.readlines()
+
+    fname_list = []
+    score_list = []
+
+    for batch_x, utt_id in data_loader:
+
+        batch_size = batch_x.size(0)
+        batch_x = batch_x.to(device)
+
+        batch_out = model(batch_x)
+
+        batch_score = (batch_out[:, 1]).data.cpu().numpy().ravel()
+        # add outputs
+        fname_list.extend(utt_id)
+        score_list.extend(batch_score.tolist())
+    assert len(trial_lines) == len(fname_list) == len(score_list)
+
+    with open(save_path, "a+") as fh:
+        for fname, cm, trl in zip(fname_list, score_list, trial_lines):
+            utt_id, key = trl.strip().split(" ")
+            assert fname == utt_id
+            fh.write("{} {} {}\n".format(fname, key, cm))
+    fh.close()
+    print("Scores saved to {}".format(save_path))
+
+
+def train_epoch(train_loader, model, lr, optim, device, args):
+    running_loss = 0
+
+    num_total = 0.0
+    algo = args.algo
+    model.train()
+    cut = 64600
+    # set objective (Loss) functions
+    weight = torch.FloatTensor([0.1, 0.9]).to(device)
+    criterion = nn.CrossEntropyLoss(weight=weight)
+    progress_bar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{args.num_epochs}")
+    for current_step, (batch_pths, batch_y) in enumerate(progress_bar):
+        batch_x = batch_pths
+        batch_size = batch_x.size(0)
+        num_total += batch_size
+
+        batch_x = batch_x.to(device)
+        batch_y = batch_y.view(-1).type(torch.int64).to(device)
+        batch_out = model(batch_x)
+
+        batch_loss = criterion(batch_out, batch_y)
+
+        running_loss += batch_loss.item() * batch_size
+
+        optimizer.zero_grad()
+        batch_loss.backward()
+        optimizer.step()
+
+    running_loss /= num_total
+
+    return running_loss
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="SSL-AASIST baseline system")
+
+    # Hyperparameters
+    parser.add_argument("--batch_size", type=int, default=64)
+    parser.add_argument("--num_epochs", type=int, default=100)
+    parser.add_argument("--lr", type=float, default=0.000001)
+    parser.add_argument("--weight_decay", type=float, default=0.0001)
+    parser.add_argument("--model_name", type=str, default="SSL-AASIST")
+    parser.add_argument("--loss", type=str, default="weighted_CCE")
+    parser.add_argument("--trn_list_path", default=None, help="path to train file")
+    parser.add_argument("--dev_list_path", default=None, help="path to validation file")
+    parser.add_argument("--test_list_path", default=None, help="path to test file")
+    parser.add_argument(
+        "--test_score_dir", default=None, help="path to save test scores"
+    )
+    # model
+    parser.add_argument(
+        "--seed", type=int, default=1234, help="random seed (default: 1234)"
+    )
+    parser.add_argument("--save_path", type=str, default=".", help="Model save path")
+    parser.add_argument("--model_path", type=str, default=None, help="Model checkpoint")
+    parser.add_argument(
+        "--comment", type=str, default=None, help="Comment to describe the saved model"
+    )
+    # Auxiliary arguments
+
+    parser.add_argument("--eval", action="store_true", default=False, help="eval mode")
+    parser.add_argument("--eval_part", type=int, default=0)
+    # backend options
+    parser.add_argument(
+        "--cudnn-deterministic-toggle",
+        action="store_false",
+        default=True,
+        help="use cudnn-deterministic? (default true)",
+    )
+
+    parser.add_argument(
+        "--cudnn-benchmark-toggle",
+        action="store_true",
+        default=False,
+        help="use cudnn-benchmark? (default false)",
+    )
+
+    ##===================================================Rawboost data augmentation ======================================================================#
+
+    parser.add_argument(
+        "--algo",
+        type=int,
+        default=5,
+        help="Rawboost algos discriptions. 0: No augmentation 1: LnL_convolutive_noise, 2: ISD_additive_noise, 3: SSI_additive_noise, 4: series algo (1+2+3), \
+                          5: series algo (1+2), 6: series algo (1+3), 7: series algo(2+3), 8: parallel algo(1,2) .[default=0]",
+    )
+
+    # LnL_convolutive_noise parameters
+    parser.add_argument(
+        "--nBands",
+        type=int,
+        default=5,
+        help="number of notch filters.The higher the number of bands, the more aggresive the distortions is.[default=5]",
+    )
+    parser.add_argument(
+        "--minF",
+        type=int,
+        default=20,
+        help="minimum centre frequency [Hz] of notch filter.[default=20] ",
+    )
+    parser.add_argument(
+        "--maxF",
+        type=int,
+        default=8000,
+        help="maximum centre frequency [Hz] (<sr/2)  of notch filter.[default=8000]",
+    )
+    parser.add_argument(
+        "--minBW",
+        type=int,
+        default=100,
+        help="minimum width [Hz] of filter.[default=100] ",
+    )
+    parser.add_argument(
+        "--maxBW",
+        type=int,
+        default=1000,
+        help="maximum width [Hz] of filter.[default=1000] ",
+    )
+    parser.add_argument(
+        "--minCoeff",
+        type=int,
+        default=10,
+        help="minimum filter coefficients. More the filter coefficients more ideal the filter slope.[default=10]",
+    )
+    parser.add_argument(
+        "--maxCoeff",
+        type=int,
+        default=100,
+        help="maximum filter coefficients. More the filter coefficients more ideal the filter slope.[default=100]",
+    )
+    parser.add_argument(
+        "--minG",
+        type=int,
+        default=0,
+        help="minimum gain factor of linear component.[default=0]",
+    )
+    parser.add_argument(
+        "--maxG",
+        type=int,
+        default=0,
+        help="maximum gain factor of linear component.[default=0]",
+    )
+    parser.add_argument(
+        "--minBiasLinNonLin",
+        type=int,
+        default=5,
+        help=" minimum gain difference between linear and non-linear components.[default=5]",
+    )
+    parser.add_argument(
+        "--maxBiasLinNonLin",
+        type=int,
+        default=20,
+        help=" maximum gain difference between linear and non-linear components.[default=20]",
+    )
+    parser.add_argument(
+        "--N_f",
+        type=int,
+        default=5,
+        help="order of the (non-)linearity where N_f=1 refers only to linear components.[default=5]",
+    )
+
+    # ISD_additive_noise parameters
+    parser.add_argument(
+        "--P",
+        type=int,
+        default=10,
+        help="Maximum number of uniformly distributed samples in [%].[defaul=10]",
+    )
+    parser.add_argument(
+        "--g_sd", type=int, default=2, help="gain parameters > 0. [default=2]"
+    )
+
+    # SSI_additive_noise parameters
+    parser.add_argument(
+        "--SNRmin",
+        type=int,
+        default=10,
+        help="Minimum SNR value for coloured additive noise.[defaul=10]",
+    )
+    parser.add_argument(
+        "--SNRmax",
+        type=int,
+        default=40,
+        help="Maximum SNR value for coloured additive noise.[defaul=40]",
+    )
+
+    ##===================================================Rawboost data augmentation ======================================================================#
+
+    load_dotenv()
+    wandb_api_key = os.getenv("WANDB_API_KEY")
+    wandb_project_name = os.getenv("WANDB_PROJECT_NAME")
+
+    if not os.path.exists("models"):
+        os.mkdir("models")
+    args = parser.parse_args()
+    wandb.login(key=wandb_api_key)
+    wandb.init(
+        project=wandb_project_name,
+        config={
+            "learning_rate": args.lr,
+            "epochs": args.num_epochs,
+            "batch_size": args.batch_size,
+            "weight_decay": args.weight_decay,
+        },
+    )
+
+    # make experiment reproducible
+    set_random_seed(args.seed, args)
+
+    # define model saving path
+    model_tag = "model_{}_{}_{}_{}".format(
+        args.loss, args.num_epochs, args.batch_size, args.lr
+    )
+    if args.comment:
+        model_tag = model_tag + "_{}".format(args.comment)
+    model_save_path = os.path.join(args.save_path, model_tag)
+
+    # set model save directory
+    if not os.path.exists(model_save_path):
+        os.mkdir(model_save_path)
+
+    # GPU device
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print("Device: {}".format(device))
+
+    model = Model(args, device)
+    nb_params = sum([param.view(-1).size()[0] for param in model.parameters()])
+    model = model.to(device)
+    print("nb_params:", nb_params)
+
+    # set Adam optimizer
+    optimizer = torch.optim.Adam(
+        model.parameters(), lr=args.lr, weight_decay=args.weight_decay
+    )
+
+    if args.model_path:
+        model.load_state_dict(torch.load(args.model_path, map_location=device))
+        print("Model loaded : {}".format(args.model_path))
+
+    # evaluation
+
+    if args.eval:
+        file_eval = genSpoof_list(
+            dir_meta=args.test_list_path, is_train=False, is_eval=True
+        )
+        print("no. of eval trials", len(file_eval))
+        eval_set = Dataset_ASVspoof2021_eval(list_IDs=file_eval)
+        eval_output = os.path.join(
+            args.test_score_dir, f"{args.model_name}_model_score.txt"
+        )
+        produce_evaluation_file(
+            eval_set, model, device, eval_output, args.test_list_path
+        )
+        output_file = os.path.join(
+            args.test_score_dir, f"{args.model_name}_model_eer.txt"
+        )
+        eval_eer = calculate_tDCF_EER(
+            cm_scores_file=eval_output, output_file=output_file
+        )
+
+        sys.exit(0)
+
+    trn_list_path = args.trn_list_path
+    dev_trial_path = args.dev_list_path
+    train_set = Dataset_ASVspoof2019_train(args, metafile=trn_list_path, algo=args.algo)
+    train_loader = DataLoader(
+        train_set,
+        batch_size=args.batch_size,
+        num_workers=16,
+        shuffle=True,
+        drop_last=True,
+    )
+    del train_set
+
+    dev_set = Dataset_ASVspoof2019_train(args, metafile=dev_trial_path, algo=args.algo)
+    dev_loader = DataLoader(
+        dev_set, batch_size=args.batch_size, num_workers=16, shuffle=False
+    )
+    del dev_set
+    # Training and validation
+    num_epochs = args.num_epochs
+    writer = SummaryWriter("logs/{}".format(model_tag))
+
+    for epoch in range(num_epochs):
+
+        running_loss = train_epoch(
+            train_loader, model, args.lr, optimizer, device, args
+        )
+        val_loss = evaluate_accuracy(dev_loader, model, device, args)
+        wandb.log({"epoch": epoch, "train_loss": running_loss, "val_loss": val_loss})
+        writer.add_scalar("val_loss", val_loss, epoch)
+        writer.add_scalar("loss", running_loss, epoch)
+        print("\n{} - {} - {} ".format(epoch, running_loss, val_loss))
+        torch.save(
+            model.state_dict(),
+            os.path.join(model_save_path, "epoch_{}.pth".format(epoch)),
+        )

audio/shiftyspeech/tests/test_api.py

@@ -0,0 +1,253 @@
+"""Tests for ShiftySpeech Audio Deepfake Detection API.
+
+Tests cover:
+- Health endpoint
+- Predict endpoint with real audio
+- Predict endpoint with fake audio
+- Audio preprocessing (pad/trim/tile)
+- Error handling for invalid input
+- Response schema validation
+"""
+
+import base64
+import io
+import os
+import sys
+import warnings
+
+import numpy as np
+import pytest
+
+warnings.filterwarnings("ignore", category=DeprecationWarning)
+
+# Monkey-patch omegaconf before importing api module
+import omegaconf._utils as _omegaconf_utils
+
+if not hasattr(_omegaconf_utils, "is_primitive_type"):
+    _omegaconf_utils.is_primitive_type = lambda t: t in (int, float, bool, str, bytes)
+
+# Add model code path for local testing
+SERVICE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+MODEL_CODE_PATH = os.path.join(
+    SERVICE_DIR, "synthetic_speech_detection", "SSL_Anti-spoofing"
+)
+if MODEL_CODE_PATH not in sys.path:
+    sys.path.insert(0, MODEL_CODE_PATH)
+
+# Patch the api module constants for local testing
+import api
+
+api.MODEL_CODE_PATH = MODEL_CODE_PATH
+api.WEIGHTS_PATH = os.path.join(SERVICE_DIR, "weights", "hfg_aug_1_2.pt")
+api.XLSR_DIR = os.path.join(SERVICE_DIR, "models")
+
+from fastapi.testclient import TestClient
+
+client = TestClient(api.app)
+
+# Dataset paths
+DATASET_DIR = os.path.join(
+    os.path.dirname(SERVICE_DIR),
+    os.pardir,
+    os.pardir,
+    "dataset",
+    "audio",
+)
+DATASET_DIR = os.path.normpath(DATASET_DIR)
+REAL_DIR = os.path.join(DATASET_DIR, "real")
+FAKE_DIR = os.path.join(DATASET_DIR, "fake")
+
+# Check if model weights are available for integration tests
+WEIGHTS_AVAILABLE = os.path.exists(api.WEIGHTS_PATH) and os.path.exists(
+    os.path.join(SERVICE_DIR, "models", "xlsr2_300m.pt")
+)
+DATASET_AVAILABLE = os.path.isdir(REAL_DIR) and os.path.isdir(FAKE_DIR)
+
+
+def _make_wav_bytes(duration_s: float = 1.0, sr: int = 16000) -> bytes:
+    """Generate a simple sine wave WAV file as bytes."""
+    import soundfile as sf
+
+    t = np.linspace(0, duration_s, int(sr * duration_s), endpoint=False)
+    audio = 0.5 * np.sin(2 * np.pi * 440 * t).astype(np.float32)
+    buf = io.BytesIO()
+    sf.write(buf, audio, sr, format="WAV")
+    buf.seek(0)
+    return buf.read()
+
+
+def _encode_file(path: str) -> str:
+    """Read a file and return base64 encoded string."""
+    with open(path, "rb") as f:
+        return base64.b64encode(f.read()).decode("utf-8")
+
+
+class TestHealthEndpoint:
+    """Tests for the /health endpoint."""
+
+    def test_health_returns_200(self):
+        response = client.get("/health")
+        assert response.status_code == 200
+
+    def test_health_contains_model_name(self):
+        response = client.get("/health")
+        data = response.json()
+        assert data["model"] == "shiftyspeech"
+
+    def test_health_contains_device(self):
+        response = client.get("/health")
+        data = response.json()
+        assert data["device"] == "cpu"
+
+    def test_health_contains_status(self):
+        response = client.get("/health")
+        data = response.json()
+        assert data["status"] in ("healthy", "degraded")
+
+
+class TestPreprocessAudio:
+    """Tests for audio preprocessing logic."""
+
+    def test_preprocess_short_audio_tiles(self):
+        """Short audio should be tiled to TARGET_SAMPLES."""
+        wav_bytes = _make_wav_bytes(duration_s=0.5, sr=16000)
+        tensor = api.preprocess_audio(wav_bytes)
+        assert tensor.shape == (1, api.TARGET_SAMPLES)
+
+    def test_preprocess_long_audio_trims(self):
+        """Long audio should be trimmed to TARGET_SAMPLES."""
+        wav_bytes = _make_wav_bytes(duration_s=10.0, sr=16000)
+        tensor = api.preprocess_audio(wav_bytes)
+        assert tensor.shape == (1, api.TARGET_SAMPLES)
+
+    def test_preprocess_exact_length(self):
+        """Audio at exact TARGET_SAMPLES should pass through."""
+        duration = api.TARGET_SAMPLES / api.SAMPLE_RATE
+        wav_bytes = _make_wav_bytes(duration_s=duration, sr=16000)
+        tensor = api.preprocess_audio(wav_bytes)
+        assert tensor.shape == (1, api.TARGET_SAMPLES)
+
+    def test_preprocess_resamples_from_8khz(self):
+        """Audio at 8kHz should be resampled to 16kHz."""
+        wav_bytes = _make_wav_bytes(duration_s=1.0, sr=8000)
+        tensor = api.preprocess_audio(wav_bytes)
+        assert tensor.shape == (1, api.TARGET_SAMPLES)
+
+    def test_preprocess_invalid_input_raises(self):
+        """Invalid audio bytes should raise ValueError."""
+        with pytest.raises(ValueError):
+            api.preprocess_audio(b"not audio data")
+
+
+@pytest.mark.skipif(
+    not WEIGHTS_AVAILABLE,
+    reason="Model weights not available locally",
+)
+class TestPredictEndpoint:
+    """Integration tests for the /predict endpoint (requires weights)."""
+
+    def test_predict_returns_200(self):
+        wav_bytes = _make_wav_bytes(duration_s=2.0)
+        b64 = base64.b64encode(wav_bytes).decode("utf-8")
+        response = client.post("/predict", json={"audio_data": b64})
+        assert response.status_code == 200
+
+    def test_predict_response_schema(self):
+        wav_bytes = _make_wav_bytes(duration_s=2.0)
+        b64 = base64.b64encode(wav_bytes).decode("utf-8")
+        response = client.post("/predict", json={"audio_data": b64})
+        data = response.json()
+        assert "model" in data
+        assert "probability" in data
+        assert "prediction" in data
+        assert "class" in data
+        assert "inference_time" in data
+        assert data["model"] == "shiftyspeech"
+
+    def test_predict_probability_in_range(self):
+        wav_bytes = _make_wav_bytes(duration_s=2.0)
+        b64 = base64.b64encode(wav_bytes).decode("utf-8")
+        response = client.post("/predict", json={"audio_data": b64})
+        data = response.json()
+        assert 0.0 <= data["probability"] <= 1.0
+
+    def test_predict_class_matches_prediction(self):
+        wav_bytes = _make_wav_bytes(duration_s=2.0)
+        b64 = base64.b64encode(wav_bytes).decode("utf-8")
+        response = client.post("/predict", json={"audio_data": b64})
+        data = response.json()
+        if data["prediction"] == 1:
+            assert data["class"] == "fake"
+        else:
+            assert data["class"] == "real"
+
+    def test_predict_custom_threshold(self):
+        wav_bytes = _make_wav_bytes(duration_s=2.0)
+        b64 = base64.b64encode(wav_bytes).decode("utf-8")
+        response = client.post(
+            "/predict",
+            json={"audio_data": b64, "threshold": 0.99},
+        )
+        data = response.json()
+        assert response.status_code == 200
+        # With threshold=0.99, only very high prob_fake => fake
+        if data["probability"] < 0.99:
+            assert data["prediction"] == 0
+            assert data["class"] == "real"
+
+    def test_predict_inference_time_positive(self):
+        wav_bytes = _make_wav_bytes(duration_s=2.0)
+        b64 = base64.b64encode(wav_bytes).decode("utf-8")
+        response = client.post("/predict", json={"audio_data": b64})
+        data = response.json()
+        assert data["inference_time"] > 0
+
+
+@pytest.mark.skipif(
+    not WEIGHTS_AVAILABLE or not DATASET_AVAILABLE,
+    reason="Model weights or dataset not available",
+)
+class TestRealDataset:
+    """Integration tests using actual dataset files."""
+
+    def test_predict_real_audio(self):
+        """Test prediction on a real audio file."""
+        path = os.path.join(REAL_DIR, "real_0.wav")
+        if not os.path.exists(path):
+            pytest.skip("real_0.wav not found")
+        b64 = _encode_file(path)
+        response = client.post("/predict", json={"audio_data": b64})
+        assert response.status_code == 200
+        data = response.json()
+        assert 0.0 <= data["probability"] <= 1.0
+
+    def test_predict_fake_audio(self):
+        """Test prediction on a fake audio file."""
+        path = os.path.join(FAKE_DIR, "fake_1.wav")
+        if not os.path.exists(path):
+            pytest.skip("fake_1.wav not found")
+        b64 = _encode_file(path)
+        response = client.post("/predict", json={"audio_data": b64})
+        assert response.status_code == 200
+        data = response.json()
+        assert 0.0 <= data["probability"] <= 1.0
+
+
+class TestPredictValidation:
+    """Tests for input validation on /predict."""
+
+    def test_predict_missing_audio_data(self):
+        response = client.post("/predict", json={})
+        assert response.status_code == 422
+
+    def test_predict_invalid_base64(self):
+        response = client.post("/predict", json={"audio_data": "not-valid-base64!!!"})
+        # Should return 500 (decode error) or 422
+        assert response.status_code in (400, 422, 500)
+
+    def test_predict_threshold_out_of_range(self):
+        response = client.post(
+            "/predict",
+            json={"audio_data": "dGVzdA==", "threshold": 1.5},
+        )
+        assert response.status_code == 422

audio/sonics/Dockerfile

@@ -0,0 +1,36 @@
+FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
+
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PYTHONUNBUFFERED=1
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    python3 python3-pip python3-dev \
+    git ffmpeg libsndfile1 \
+    build-essential g++ \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN ln -sf /usr/bin/python3 /usr/bin/python
+
+WORKDIR /app
+
+# Install PyTorch with CUDA 12.1
+RUN pip install --no-cache-dir \
+    torch==2.5.1 torchaudio==2.5.1 \
+    --index-url https://download.pytorch.org/whl/cu121
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY app.py .
+
+# Pre-download model weights
+RUN python -c "from sonics import HFAudioClassifier; HFAudioClassifier.from_pretrained('awsaf49/sonics-spectttra-alpha-120s')"
+
+EXPOSE 8003
+
+RUN adduser --disabled-password --gecos '' appuser
+USER appuser
+
+ENV PRELOAD_MODEL=true
+
+CMD ["python", "app.py"]

audio/sonics/app.py

@@ -0,0 +1,285 @@
+"""SONICS (SpecTTTra) Synthetic Music Detection API.
+
+Detects AI-generated music (Suno, Udio, etc.) using the SpecTTTra
+architecture from the SONICS project (ICLR 2025).
+
+The model performs binary classification on raw audio waveforms via
+internal MelSpectrogram features. It outputs a single logit; we apply
+sigmoid to obtain the fake probability.
+
+Reference: https://github.com/awsaf49/sonics
+"""
+
+import base64
+import io
+import logging
+import os
+import platform
+import time
+from typing import Optional
+
+import librosa
+import numpy as np
+import torch
+import uvicorn
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger("sonics_api")
+
+# Constants
+MODEL_NAME = "sonics_detection"
+HF_MODEL_ID = "awsaf49/sonics-spectttra-alpha-120s"
+SAMPLE_RATE = 16000
+MAX_TIME = 120  # seconds (matches alpha-120s config)
+MAX_LEN = MAX_TIME * SAMPLE_RATE  # 1_920_000 samples
+PRELOAD_MODEL = os.environ.get("PRELOAD_MODEL", "true").lower() == "true"
+MODEL_TIMEOUT = int(os.environ.get("MODEL_TIMEOUT", "300"))
+
+
+def _get_device():
+    """Select optimal device: MPS (Apple) > CUDA (NVIDIA) > CPU."""
+    override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower()
+    if override == "cpu":
+        return torch.device("cpu")
+    if override == "cuda" and torch.cuda.is_available():
+        return torch.device("cuda")
+    if (
+        override == "mps"
+        and hasattr(torch.backends, "mps")
+        and torch.backends.mps.is_available()
+    ):
+        return torch.device("mps")
+    if override:
+        pass  # Invalid override, fall through to auto-detect
+    if (
+        platform.system() == "Darwin"
+        and hasattr(torch.backends, "mps")
+        and torch.backends.mps.is_available()
+    ):
+        return torch.device("mps")
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    return torch.device("cpu")
+
+
+DEVICE = _get_device()
+
+if DEVICE.type == "cuda":
+    torch.backends.cudnn.benchmark = True
+    torch.set_float32_matmul_precision("high")
+
+if DEVICE.type == "cuda":
+    logger.info(
+        "Device: cuda (%s, %.1f GB VRAM)",
+        torch.cuda.get_device_name(0),
+        torch.cuda.get_device_properties(0).total_memory / 1024**3,
+    )
+else:
+    logger.warning(
+        "Device: %s (no CUDA available -- check nvidia-container-toolkit)",
+        DEVICE,
+    )
+
+# Global model instance
+model = None
+
+
+class AudioInput(BaseModel):
+    """Request schema for audio deepfake detection."""
+
+    audio_data: str = Field(
+        ..., description="Base64 encoded audio string (WAV/MP3/FLAC/etc)"
+    )
+    threshold: Optional[float] = Field(
+        0.5, ge=0.0, le=1.0, description="Classification threshold"
+    )
+
+
+app = FastAPI(
+    title="SONICS Synthetic Music Detection API",
+    description=(
+        "Service for detecting AI-generated music using the "
+        "SpecTTTra model from the SONICS project (ICLR 2025)."
+    ),
+    version="1.0.0",
+)
+
+
+def load_model():
+    """Load the SONICS HFAudioClassifier from HuggingFace Hub.
+
+    Returns:
+        The loaded model, or None if loading fails.
+    """
+    global model
+    if model is not None:
+        return model
+
+    logger.info("Loading SONICS model '%s' onto %s...", HF_MODEL_ID, DEVICE)
+
+    try:
+        from sonics import HFAudioClassifier
+
+        model = HFAudioClassifier.from_pretrained(
+            HF_MODEL_ID,
+            map_location=str(DEVICE),
+        )
+        model.to(DEVICE)
+        model.eval()
+
+        logger.info("SONICS model loaded successfully.")
+        return model
+    except Exception:
+        logger.exception("Failed to load SONICS model")
+        model = None
+        return None
+
+
+@app.on_event("startup")
+async def startup_event():
+    """Optionally preload model on service startup."""
+    if PRELOAD_MODEL:
+        load_model()
+
+
+@app.get("/")
+async def root():
+    """Root info endpoint."""
+    return {
+        "service": "SONICS Synthetic Music Detection",
+        "model": MODEL_NAME,
+        "version": "1.0.0",
+    }
+
+
+def _gpu_health_info() -> dict:
+    """Return GPU metrics for the health endpoint."""
+    if torch.cuda.is_available() and DEVICE.type == "cuda":
+        return {
+            "gpu_name": torch.cuda.get_device_name(0),
+            "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2),
+            "vram_total_mb": round(
+                torch.cuda.get_device_properties(0).total_memory / 1024**2
+            ),
+        }
+    return {}
+
+
+@app.get("/health")
+async def health():
+    """Health check endpoint."""
+    return {
+        "status": "healthy" if model is not None else "degraded",
+        "model": MODEL_NAME,
+        "device": str(DEVICE),
+        **_gpu_health_info(),
+    }
+
+
+def preprocess_audio(audio_bytes: bytes) -> torch.Tensor:
+    """Preprocess audio for SONICS inference.
+
+    Loads audio from raw bytes, resamples to 16 kHz mono,
+    crops or zero-pads to MAX_LEN samples, and normalises by
+    standard deviation (matching the training pipeline).
+
+    Args:
+        audio_bytes: Raw audio file bytes (WAV, MP3, FLAC, etc.).
+
+    Returns:
+        Audio tensor of shape (1, MAX_LEN) on DEVICE.
+
+    Raises:
+        ValueError: If audio preprocessing fails.
+    """
+    try:
+        logger.info("Starting audio preprocessing...")
+        audio, sr = librosa.load(io.BytesIO(audio_bytes), sr=SAMPLE_RATE, mono=True)
+        logger.info("Audio loaded. Length: %d samples at %dHz", len(audio), sr)
+
+        # Crop or pad to fixed length (matching SONICS dataset.py)
+        if len(audio) > MAX_LEN:
+            # Crop from 3/4 position (matching eval-mode logic)
+            idx = int((len(audio) - MAX_LEN) / 4 * 3)
+            audio = audio[idx : idx + MAX_LEN]
+        elif len(audio) < MAX_LEN:
+            audio = np.pad(audio, (0, MAX_LEN - len(audio)), mode="constant")
+
+        # Normalise by standard deviation (matching training pipeline)
+        audio /= np.maximum(np.std(audio), 1e-6)
+
+        logger.info("Audio preprocessed to %d samples", len(audio))
+
+        audio_tensor = torch.from_numpy(audio).float().unsqueeze(0)
+        audio_tensor = audio_tensor.to(DEVICE)
+        return audio_tensor
+    except Exception as e:
+        logger.error("Error preprocessing audio: %s", e)
+        raise ValueError(f"Audio preprocessing failed: {str(e)}")
+
+
+@app.post("/predict")
+async def predict(input_data: AudioInput):
+    """Run synthetic music detection on base64-encoded audio.
+
+    The model uses BCEWithLogitsLoss with num_classes=1, so it
+    outputs a single logit.  We apply sigmoid to obtain the fake
+    probability.
+    """
+    if model is None:
+        if load_model() is None:
+            raise HTTPException(status_code=503, detail="Model not loaded")
+
+    try:
+        start_time = time.time()
+        logger.info(
+            "Prediction request. Data size: %d chars",
+            len(input_data.audio_data),
+        )
+
+        # Decode base64 audio
+        audio_bytes = base64.b64decode(input_data.audio_data)
+
+        # Preprocess
+        audio_tensor = preprocess_audio(audio_bytes)
+
+        # Inference
+        logger.info("Starting model inference...")
+        with torch.no_grad():
+            logits = model(audio_tensor)
+            # logits shape: (1, 1) -- single logit for binary classification
+            prob_fake = torch.sigmoid(logits).squeeze().item()
+
+        prediction = 1 if prob_fake >= input_data.threshold else 0
+        verdict = "fake" if prediction == 1 else "real"
+        inference_time = time.time() - start_time
+
+        logger.info(
+            "Prediction: %s (prob_fake=%.4f, time=%.3fs)",
+            verdict,
+            prob_fake,
+            inference_time,
+        )
+
+        return {
+            "model": MODEL_NAME,
+            "probability": float(prob_fake),
+            "prediction": int(prediction),
+            "class": verdict,
+            "inference_time": float(inference_time),
+        }
+
+    except Exception as e:
+        logger.exception("Error during prediction: %s", e)
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+if __name__ == "__main__":
+    port = int(os.environ.get("MODEL_PORT", 8003))
+    uvicorn.run(app, host="0.0.0.0", port=port)

audio/sonics/requirements.txt

@@ -0,0 +1,22 @@
+# SONICS model (pip-installable from GitHub)
+sonics @ git+https://github.com/awsaf49/sonics.git
+
+# Core inference
+torch>=2.4.0
+torchaudio>=2.4.0
+
+# Audio processing
+librosa>=0.9.0
+
+# ML utilities (SONICS dependency)
+timm>=1.0.7
+fvcore
+
+# HuggingFace Hub (for model download)
+huggingface-hub
+
+# API
+fastapi
+uvicorn[standard]
+pydantic>=2.0
+numpy

ensemble-core/Dockerfile

@@ -0,0 +1,10 @@
+FROM python:3.9-slim
+WORKDIR /app
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+COPY main.py .
+# Drop root privileges
+RUN adduser --disabled-password --gecos '' appuser
+USER appuser
+
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8003"]

ensemble-core/main.py

@@ -0,0 +1,23 @@
+from typing import Any, Dict
+
+from fastapi import FastAPI
+from pydantic import BaseModel
+
+app = FastAPI(title="Ensemble Core Service")
+
+
+class EnsembleRequest(BaseModel):
+    media_type: str
+    model_results: Dict[str, Any]
+    method: str = "stacking"
+
+
+@app.post("/calculate")
+async def calculate_ensemble(request: EnsembleRequest):
+    # TODO: Migrate proprietary ensemble logic from gateway
+    return {"verdict": "fake", "confidence": 0.95, "method_used": request.method}
+
+
+@app.get("/health")
+async def health():
+    return {"status": "healthy"}

ensemble-core/requirements.txt

@@ -0,0 +1 @@
+fastapi\nuvicorn\npydantic

ensemble-core/scripts/create_dataset.py

@@ -0,0 +1,169 @@
+#!/usr/bin/env python3
+"""
+Build a balanced 10k-real / 10k-fake face-image folder with maximum
+deepfake-tech variety.
+
+Usage:
+    python build_face_dataset.py --out_dir ./faces20k --seed 42
+"""
+
+import argparse
+import os
+import pathlib
+import random
+import shutil
+import subprocess
+import sys
+import zipfile
+from collections import defaultdict
+
+import pandas as pd
+from tqdm import tqdm
+
+# ----------------------------------------------------------------------
+# 1. Edit here to add / remove sources
+# ----------------------------------------------------------------------
+DATASETS = [
+    {
+        "name": "140k",
+        "slug": "xhlulu/140k-real-and-fake-faces",
+        "subdirs": {"real": "real", "fake": "fake"},
+        "fake_label": "stylegan2",
+    },
+    {
+        "name": "deepfake_real",
+        "slug": "manjilkarki/deepfake-and-real-images",
+        "subdirs": {"real": "real", "fake": "fake"},
+        "fake_label": "pggan_stylegan_mix",
+    },
+    {
+        "name": "dfdc_f150",
+        "slug": "sciarrilli/dfdc-f150",
+        "subdirs": {"real": "real", "fake": "fake"},
+        "fake_label": "dfdc_swaps",
+    },
+    {
+        "name": "faceforensics_imgs",
+        "slug": "greatgamedota/faceforensics",
+        "subdirs": {"real": "real", "fake": "fake"},
+        "fake_label": "ffpp_swaps",
+    },
+]
+
+TARGET_PER_CLASS = 10_000
+# ----------------------------------------------------------------------
+
+
+def kaggle_download(slug: str, dest: pathlib.Path) -> pathlib.Path:
+    """Download <slug> to dest/. Returns path of the zip."""
+    dest.mkdir(parents=True, exist_ok=True)
+    zip_path = dest / f"{slug.split('/')[-1]}.zip"
+    if zip_path.exists():
+        return zip_path
+    print(f"Downloading {slug} …")
+    subprocess.run(
+        ["kaggle", "datasets", "download", "-d", slug, "-p", str(dest), "--quiet"],
+        check=True,
+    )
+    return zip_path
+
+
+def extract(zip_path: pathlib.Path, dest: pathlib.Path) -> pathlib.Path:
+    """Unzip if needed. Returns extraction dir."""
+    extract_dir = dest / zip_path.stem
+    if extract_dir.exists():
+        return extract_dir
+    print(f"Extracting {zip_path.name} …")
+    with zipfile.ZipFile(zip_path) as zf:
+        zf.extractall(path=extract_dir)
+    return extract_dir
+
+
+def glob_images(root: pathlib.Path, pattern: str):
+    return list(root.glob(pattern)) + list(root.glob(pattern.replace("jpg", "png")))
+
+
+def main(out_dir: pathlib.Path, seed: int):
+    random.seed(seed)
+    temp_root = out_dir / "_raw"
+    real_pool, fake_pool = [], []
+    fake_source_tag = {}  # path -> dataset tag
+
+    # ------------------------------------------------------------------
+    # 2. Pull sources
+    # ------------------------------------------------------------------
+    for ds in DATASETS:
+        zip_path = kaggle_download(ds["slug"], temp_root)
+        extract_dir = extract(zip_path, temp_root)
+        real_dir = extract_dir / ds["subdirs"]["real"]
+        fake_dir = extract_dir / ds["subdirs"]["fake"]
+        real_pool += glob_images(real_dir, "**/*.jpg")
+        fakes = glob_images(fake_dir, "**/*.jpg")
+        fake_pool += fakes
+        for fp in fakes:
+            fake_source_tag[str(fp)] = ds["fake_label"]
+
+    # sanity check
+    if len(real_pool) < TARGET_PER_CLASS or len(fake_pool) < TARGET_PER_CLASS:
+        print("Not enough images – add another dataset.", file=sys.stderr)
+        sys.exit(1)
+
+    # ------------------------------------------------------------------
+    # 3. Sample
+    # ------------------------------------------------------------------
+    random.shuffle(real_pool)
+    random.shuffle(fake_pool)
+
+    # try to spread fake quota equally over sources
+    per_source_quota = TARGET_PER_CLASS // len(DATASETS)
+    selected_fake = []
+    taken = defaultdict(int)
+    for fp in fake_pool:
+        tag = fake_source_tag[str(fp)]
+        if taken[tag] < per_source_quota:
+            selected_fake.append(fp)
+            taken[tag] += 1
+        if len(selected_fake) == TARGET_PER_CLASS:
+            break
+    # top-up if we’re short (some sets too small)
+    if len(selected_fake) < TARGET_PER_CLASS:
+        needed = TARGET_PER_CLASS - len(selected_fake)
+        selected_fake += fake_pool[len(selected_fake) : len(selected_fake) + needed]
+
+    selected_real = real_pool[:TARGET_PER_CLASS]
+
+    # ------------------------------------------------------------------
+    # 4. Copy to final tree + manifest
+    # ------------------------------------------------------------------
+    for cls in ("real", "fake"):
+        (out_dir / cls).mkdir(parents=True, exist_ok=True)
+
+    manifest_rows = []
+
+    def copy_files(file_list, cls):
+        for src in tqdm(file_list, desc=f"Copying {cls}"):
+            dst = out_dir / cls / src.name
+            shutil.copy(src, dst)
+            manifest_rows.append(
+                {
+                    "filename": dst.name,
+                    "label": cls,
+                    "source": fake_source_tag.get(str(src), "n/a"),
+                }
+            )
+
+    copy_files(selected_real, "real")
+    copy_files(selected_fake, "fake")
+
+    pd.DataFrame(manifest_rows).to_csv(out_dir / "manifest.csv", index=False)
+    print("Done →", out_dir)
+
+
+if __name__ == "__main__":
+    p = argparse.ArgumentParser()
+    p.add_argument(
+        "--out_dir", default="faces20k", type=pathlib.Path, help="destination folder"
+    )
+    p.add_argument("--seed", default=42, type=int)
+    args = p.parse_args()
+    main(args.out_dir, args.seed)

ensemble-core/scripts/meta_feature_generator.py

@@ -0,0 +1,366 @@
+#!/usr/bin/env python3
+"""
+DeepSafe Meta-Feature Generator
+===============================
+
+Orchestrates the generation of meta-feature datasets for training stacking ensembles.
+This component acts as a data ingestion pipeline that:
+1.  Scans a target directory for labeled media (Real/Fake).
+2.  Queries the distributed model microservices to obtain base probability scores.
+3.  Aggregates these scores into a structured feature matrix (CSV) for the meta-learner.
+
+Architectural Note:
+This script is designed to be fault-tolerant. If a specific model microservice is unreachable
+or fails for a subset of files, the pipeline continues, recording NaNs for those features.
+This ensures that a single model failure does not halt the entire training data generation process,
+though downstream imputers must handle these missing values.
+"""
+
+import argparse
+import gc
+import json
+import os
+import sys
+import time
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+import pandas as pd
+from rich.console import Console
+from rich.panel import Panel
+from rich.progress import (
+    BarColumn,
+    MofNCompleteColumn,
+    Progress,
+    SpinnerColumn,
+    TextColumn,
+    TimeElapsedColumn,
+)
+from rich.table import Table
+
+# Ensure utils is importable regardless of execution context.
+# This fallback is necessary when running the script directly from the project root
+# without an installed package structure.
+try:
+    from utils.api_client import APIClient
+    from utils.config_manager import ConfigManager
+    from utils.media_handler import MediaHandler
+except ImportError:
+    project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))
+    if project_root not in sys.path:
+        sys.path.insert(0, project_root)
+    try:
+        from utils.api_client import APIClient
+        from utils.config_manager import ConfigManager
+        from utils.media_handler import MediaHandler
+    except ImportError as e:
+        print(f"Critical Error: Failed to resolve utils dependency. {e}")
+        sys.exit(1)
+
+
+console = Console(width=120)
+
+
+class MetaFeatureGenerator:
+    """
+    Manages the ETL process for meta-learning datasets.
+
+    Attributes:
+        media_type (str): The domain of operation (image, video, audio).
+        config_manager (ConfigManager): Centralized configuration handler.
+        api_client (APIClient): Interface for communicating with model microservices.
+    """
+
+    def __init__(self, media_type: str, config_manager: ConfigManager):
+        self.media_type = media_type
+        self.config_manager = config_manager
+        # run_from_host=True implies we are running outside the docker network (e.g., local dev),
+        # so we use localhost ports mapped in docker-compose.
+        self.api_client = APIClient(config_manager, media_type, run_from_host=True)
+        self.media_handler = MediaHandler(config_manager)
+        self.base_model_names = list(
+            config_manager.get_model_endpoints(media_type).keys()
+        )
+
+        if not self.base_model_names:
+            console.print(
+                f"[bold red]Configuration Error: No base models defined for '{media_type}'.[/bold red]"
+            )
+            sys.exit(1)
+
+    def generate(
+        self,
+        input_dir: str,
+        output_csv_path: str,
+        default_threshold: float,
+        specific_models: Optional[List[str]] = None,
+    ):
+        """
+        Executes the generation pipeline.
+
+        Args:
+            input_dir: Root directory containing 'Real' and 'Fake' subdirectories.
+            output_csv_path: Destination for the resulting feature matrix.
+            default_threshold: Decision threshold passed to models (mostly for logging/reference).
+            specific_models: Optional filter to run only a subset of available models.
+        """
+
+        console.print(
+            Panel(
+                f"[bold cyan]Meta-Feature Generation Protocol ({self.media_type.capitalize()})[/bold cyan]\n"
+                f"Source: {input_dir}\n"
+                f"Target: {output_csv_path}\n"
+                f"Active Models: {specific_models or 'All configured'}",
+                title="Pipeline Configuration",
+                border_style="blue",
+                expand=False,
+            )
+        )
+
+        # Discovery phase: Scan filesystem for valid media files and infer ground truth from directory structure.
+        media_files_with_gt = self.media_handler.find_media_files(
+            input_dir, self.media_type
+        )
+        if not media_files_with_gt:
+            console.print(
+                f"[bold red]Abort: No valid {self.media_type} files found in '{input_dir}'.[/bold red]"
+            )
+            return
+
+        # Determine the execution scope (subset of models vs all).
+        models_to_query = self.base_model_names
+        if specific_models:
+            models_to_query = [m for m in specific_models if m in self.base_model_names]
+            if not models_to_query:
+                console.print(
+                    f"[bold red]Configuration Mismatch: Requested models {specific_models} are not configured for '{self.media_type}'.[/bold red]"
+                )
+                return
+            console.print(f"Scope restricted to: {models_to_query}")
+
+        all_feature_data = []
+
+        # Execution phase: Iterate through files and query models.
+        # We use a rich progress bar for observability during long-running batch processes.
+        with Progress(
+            SpinnerColumn(),
+            TextColumn("[progress.description]{task.description}"),
+            BarColumn(),
+            MofNCompleteColumn(),
+            TimeElapsedColumn(),
+        ) as progress:
+            total_files = len(media_files_with_gt)
+            outer_task = progress.add_task(
+                f"Processing {self.media_type} corpus...", total=total_files
+            )
+
+            for file_idx, (media_path, ground_truth_label) in enumerate(
+                media_files_with_gt
+            ):
+                media_file_name = os.path.basename(media_path)
+                progress.update(
+                    outer_task,
+                    description=f"Processing: [cyan]{media_file_name}[/cyan]",
+                )
+
+                # Pre-encode media to base64 once to avoid redundant I/O operations per model.
+                encoded_media = self.media_handler.encode_media_to_base64(media_path)
+                if not encoded_media:
+                    console.print(
+                        f"[yellow]Skip: Encoding failed for {media_file_name}.[/yellow]"
+                    )
+                    progress.advance(outer_task)
+                    continue
+
+                # Feature vector initialization
+                current_media_features: Dict[str, Any] = {
+                    "media_path": media_path,
+                    "media_name": media_file_name,
+                    # Map string labels to numeric binary targets: Fake=1, Real=0.
+                    "ground_truth": (
+                        1
+                        if ground_truth_label == "Fake"
+                        else (0 if ground_truth_label == "Real" else -1)
+                    ),
+                }
+
+                # Initialize feature columns with NaN. This ensures structural consistency in the DataFrame
+                # even if specific model queries fail.
+                for model_name_cfg in self.base_model_names:
+                    current_media_features[f"{model_name_cfg}_prob"] = np.nan
+
+                # Query loop
+                for model_name_query in models_to_query:
+                    model_result = self.api_client.test_with_individual_model(
+                        model_name_query, media_path, encoded_media, default_threshold
+                    )
+
+                    if (
+                        "error" not in model_result
+                        and model_result.get("probability") is not None
+                    ):
+                        current_media_features[f"{model_name_query}_prob"] = (
+                            model_result["probability"]
+                        )
+                    else:
+                        # Log failure but do not interrupt the pipeline. Robustness is key here.
+                        error_msg = model_result.get(
+                            "error", "Invalid response payload"
+                        )
+                        console.print(
+                            f"[yellow]Model Failure: {model_name_query} on {media_file_name}. Reason: {error_msg}.[/yellow]",
+                            highlight=False,
+                        )
+
+                all_feature_data.append(current_media_features)
+                progress.advance(outer_task)
+
+                # Explicit garbage collection to prevent memory bloat during large dataset processing.
+                gc.collect()
+
+        if not all_feature_data:
+            console.print(
+                "[bold red]Pipeline Failure: No features generated.[/bold red]"
+            )
+            return
+
+        # Data serialization and validation
+        meta_features_df = pd.DataFrame(all_feature_data)
+
+        # Filter invalid ground truth (should be handled by discovery, but defensive programming is good).
+        meta_features_df = meta_features_df[meta_features_df["ground_truth"] != -1]
+
+        if meta_features_df.empty:
+            console.print(
+                "[bold red]Data Error: No valid labeled data remaining after processing.[/bold red]"
+            )
+            return
+
+        # Schema enforcement: Ensure all expected columns exist.
+        expected_prob_cols = [f"{mn}_prob" for mn in self.base_model_names]
+        for col in expected_prob_cols:
+            if col not in meta_features_df.columns:
+                meta_features_df[col] = np.nan
+
+        # Column ordering for readability and consistency.
+        ordered_prob_cols = sorted(
+            [col for col in meta_features_df.columns if col.endswith("_prob")]
+        )
+        final_cols_order = (
+            ["media_path", "media_name"] + ordered_prob_cols + ["ground_truth"]
+        )
+        meta_features_df = meta_features_df[final_cols_order]
+
+        try:
+            os.makedirs(
+                os.path.dirname(os.path.abspath(output_csv_path)), exist_ok=True
+            )
+            meta_features_df.to_csv(output_csv_path, index=False, float_format="%.6f")
+
+            console.print(
+                f"\n[bold green]Success: Dataset persisted to {os.path.abspath(output_csv_path)}[/bold green]"
+            )
+            console.print(f"Dimensions: {meta_features_df.shape}")
+
+            # Quality Assurance: Report missing values to inform downstream handling strategies.
+            nan_summary_table = Table(
+                title="Data Quality Report (Missing Values)", show_lines=True
+            )
+            nan_summary_table.add_column("Feature", style="cyan")
+            nan_summary_table.add_column(
+                "Missing Count", style="magenta", justify="right"
+            )
+            nan_summary_table.add_column("Missing %", style="yellow", justify="right")
+
+            for col in ordered_prob_cols:
+                nan_count = meta_features_df[col].isnull().sum()
+                nan_percent = (
+                    (nan_count / len(meta_features_df)) * 100
+                    if len(meta_features_df) > 0
+                    else 0
+                )
+                nan_summary_table.add_row(col, str(nan_count), f"{nan_percent:.2f}%")
+            console.print(nan_summary_table)
+
+        except Exception as e:
+            console.print(
+                f"[bold red]I/O Error: Failed to write output CSV. {e}[/bold red]"
+            )
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="DeepSafe Meta-Feature Generator: ETL for Stacking Ensemble Training Data.",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument(
+        "--media-type",
+        type=str,
+        choices=["image", "video", "audio"],
+        required=True,
+        help="Target domain. Defines the model registry subset.",
+    )
+    parser.add_argument(
+        "--input-dir",
+        type=str,
+        required=True,
+        help="Source directory. Must contain 'Real' and 'Fake' subdirectories for label inference.",
+    )
+    parser.add_argument(
+        "--output-csv",
+        type=str,
+        required=True,
+        help="Destination path for the generated feature matrix.",
+    )
+    parser.add_argument(
+        "--threshold",
+        type=float,
+        help="Decision threshold override (0.0-1.0). Defaults to system config.",
+    )
+    parser.add_argument(
+        "--specific-models",
+        type=str,
+        help="Optional filter: Comma-separated list of model identifiers to query.",
+    )
+    parser.add_argument(
+        "--config-path", type=str, default=None, help=f"Configuration override path."
+    )
+
+    args = parser.parse_args()
+
+    # Initialize configuration subsystem
+    cfg_manager = ConfigManager(config_path=args.config_path)
+    if not cfg_manager.is_config_loaded_successfully():
+        sys.exit(1)
+
+    default_thresh_from_config = cfg_manager.get_default("default_threshold", 0.5)
+    query_threshold = (
+        args.threshold if args.threshold is not None else default_thresh_from_config
+    )
+
+    specific_models_list = (
+        [m.strip() for m in args.specific_models.split(",")]
+        if args.specific_models
+        else None
+    )
+
+    generator = MetaFeatureGenerator(args.media_type, cfg_manager)
+    generator.generate(
+        args.input_dir, args.output_csv, query_threshold, specific_models_list
+    )
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except KeyboardInterrupt:
+        console.print("\n[bold yellow]Process Interrupted by User.[/bold yellow]")
+        sys.exit(0)
+    except Exception as e:
+        console.print(f"\n[bold red]Fatal Error: {e}[/bold red]")
+        import traceback
+
+        console.print(
+            Panel(traceback.format_exc(), title="Stack Trace", border_style="red")
+        )
+        sys.exit(1)

ensemble-core/scripts/train_meta_learner_advanced.py

@@ -0,0 +1,1228 @@
+#!/usr/bin/env python3
+"""
+DeepSafe Advanced Meta-Learner Training Suite (train_meta_learner_advanced.py)
+==============================================================================
+
+This script trains and evaluates various meta-learners (stacking ensembles)
+for deepfake detection. It takes a CSV file of meta-features (outputs from
+base deepfake detection models) and ground truth labels as input.
+
+Key Features:
+-------------
+1.  Modality-Specific Training: Supports training separate meta-learners for
+    different media types (image, video, audio) using the `--media-type` argument.
+    This ensures that the meta-learner is optimized for the characteristics of
+    the base models relevant to that modality.
+2.  Data Preprocessing: Includes imputation for missing values (e.g., if a base
+    model failed) and feature scaling.
+3.  Multiple Meta-Learner Models: Trains and evaluates several standard classifiers
+    (Logistic Regression, Random Forest, Gradient Boosting, SVC, KNN, Naive Bayes)
+    and, if available, advanced models like XGBoost and LightGBM.
+4.  Hyperparameter Optimization:
+    - Supports Optuna for efficient hyperparameter search.
+    - Falls back to GridSearchCV if Optuna is not installed or if specified.
+5.  Comprehensive Evaluation:
+    - Calculates Accuracy, F1-Score, Precision, Recall, and ROC AUC for each model.
+    - Generates classification reports and confusion matrices.
+    - Plots ROC curves for visual comparison of all trained meta-learners and
+      simple ensemble baselines.
+6.  Simple Ensemble Baselines: Also evaluates simple averaging and majority vote
+    ensembles for comparison against more complex stacking models. Includes an
+    option for optimized weighted averaging.
+7.  Artifact Generation:
+    - Saves all trained meta-learner models (e.g., .joblib files).
+    - Saves the data preprocessor (imputer + scaler).
+    - Saves the list of feature columns used during training.
+    - Saves a summary of all experiment metrics in JSON format.
+    - The final, best-performing trainable meta-learner and its associated
+      preprocessors are saved with generic names inside media-type specific
+      subfolders (e.g., api_artifacts_dir/image/deepsafe_meta_learner.joblib).
+8.  Configurable Output: Allows specifying separate directories for general
+    experiment outputs and for API-ready deployment artifacts.
+
+CLI Usage:
+----------
+python train_meta_learner_advanced.py \\
+    --media-type [image|video|audio] \\
+    --meta-file /path/to/meta_features_[media_type].csv \\
+    --output-dir ./meta_learning_experiment_runs/ \\
+    --api-artifacts-dir ./api/meta_model_artifacts/ \\
+    [--optimizer optuna|gridsearch] \\
+    [--optuna-trials 50] \\
+    [--weights /path/to/custom_weights.json]
+
+Arguments:
+----------
+  --media-type {image,video,audio}
+                        (Required) The type of media for which the meta-learner
+                        is being trained. This affects output artifact naming.
+  --meta-file META_FILE
+                        (Required) Path to the CSV file containing meta-features
+                        (base model outputs) and a 'ground_truth' column.
+  --output-dir OUTPUT_DIR
+                        Base directory for saving all experiment-related outputs
+                        (logs, plots, individual model files from this run).
+                        A timestamped, media-type-specific subdirectory will be created.
+                        (Default: ./meta_learning_experiment_runs/)
+  --api-artifacts-dir API_ARTIFACTS_DIR
+                        Directory to save the final, API-ready deployment artifacts
+                        (e.g., ./api/meta_model_artifacts/image/deepsafe_meta_learner.joblib).
+                        (Default: ./api/meta_model_artifacts/)
+  --optimizer {optuna,gridsearch}
+                        Hyperparameter optimization strategy (Default: optuna).
+  --optuna-trials N
+                        Number of trials for Optuna optimization (Default: 50).
+  --weights WEIGHTS_PATH_OR_JSON
+                        Optional. Path to a JSON file or a JSON string defining
+                        custom weights for the 'Provided_Weighted_Average' ensemble.
+                        Keys should be base model names (without '_prob' suffix).
+
+Example (Image Meta-Learner):
+-----------------------------
+python train_meta_learner_advanced.py \\
+    --media-type image \\
+    --meta-file ./meta_learning_data/meta_features_image.csv \\
+    --output-dir ./ml_experiments_images \\
+    --api-artifacts-dir ./deepsafe_private/api/meta_model_artifacts \\
+    --optimizer optuna \\
+    --optuna-trials 100
+
+This will train image-specific meta-learners, save experiment details in
+`./ml_experiments_images/experiments_image_YYYYMMDD_HHMMSS/`, and place
+API-ready artifacts like `deepsafe_meta_learner.joblib` into
+`./deepsafe_private/api/meta_model_artifacts/image/`.
+"""
+
+import argparse
+import itertools
+import json
+import os
+import time
+from typing import Any, Dict, List, Optional, Tuple
+
+import joblib
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import seaborn as sns
+from rich.console import Console
+from rich.panel import Panel
+from rich.progress import (
+    BarColumn,
+    MofNCompleteColumn,
+    Progress,
+    SpinnerColumn,
+    TextColumn,
+    TimeElapsedColumn,
+)
+from rich.table import Table
+from sklearn.ensemble import GradientBoostingClassifier, RandomForestClassifier
+from sklearn.impute import SimpleImputer
+from sklearn.linear_model import LogisticRegression
+from sklearn.metrics import (
+    accuracy_score,
+    auc,
+    classification_report,
+    confusion_matrix,
+    f1_score,
+    precision_score,
+    recall_score,
+    roc_auc_score,
+    roc_curve,
+)
+from sklearn.model_selection import StratifiedKFold, train_test_split
+from sklearn.naive_bayes import GaussianNB
+from sklearn.neighbors import KNeighborsClassifier
+from sklearn.pipeline import Pipeline
+from sklearn.preprocessing import StandardScaler
+from sklearn.svm import SVC
+
+# --- Optional Advanced Hyperparameter Optimization & Models ---
+OPTIMIZER_CHOICE_DEFAULT = "optuna"
+
+try:
+    import optuna
+
+    OPTIMIZER_AVAILABLE_OPTUNA = True
+except ImportError:
+    optuna = None
+    OPTIMIZER_AVAILABLE_OPTUNA = False
+
+from sklearn.model_selection import GridSearchCV
+
+try:
+    from xgboost import XGBClassifier
+
+    XGBOOST_AVAILABLE = True
+except ImportError:
+    XGBClassifier = None
+    XGBOOST_AVAILABLE = False
+
+try:
+    from lightgbm import LGBMClassifier
+
+    LIGHTGBM_AVAILABLE = True
+except ImportError:
+    LGBMClassifier = None
+    LIGHTGBM_AVAILABLE = False
+
+console = Console(width=120)
+
+# --- Configuration ---
+DEFAULT_EXPERIMENT_OUTPUT_DIR_BASE = "./meta_learning_experiment_runs"
+DEFAULT_API_ARTIFACTS_DIR = "./api/meta_model_artifacts"
+DEFAULT_THRESHOLD_FOR_SIMPLE_ENSEMBLES = 0.5
+N_OPTUNA_TRIALS_DEFAULT = 50
+CV_FOLDS_DEFAULT = 5
+
+
+# --- Helper Functions ---
+class NpEncoder(json.JSONEncoder):
+    def default(self, o: Any) -> Any:
+        if isinstance(o, np.integer):
+            return int(o)
+        if isinstance(o, np.floating):
+            return float(o)
+        if isinstance(o, np.ndarray):
+            return o.tolist()
+        return super(NpEncoder, self).default(o)
+
+
+def evaluate_model_predictions(
+    y_true: np.ndarray,
+    y_pred_class: np.ndarray,
+    y_pred_proba: Optional[np.ndarray],
+    model_name: str = "Model",
+) -> Dict[str, Any]:
+    metrics: Dict[str, Any] = {"name": model_name}
+    try:
+        metrics["accuracy"] = accuracy_score(y_true, y_pred_class)
+        metrics["f1_score"] = f1_score(y_true, y_pred_class, zero_division=0)
+        metrics["precision"] = precision_score(y_true, y_pred_class, zero_division=0)
+        metrics["recall"] = recall_score(y_true, y_pred_class, zero_division=0)
+
+        roc_auc_val = np.nan
+        if y_pred_proba is not None and len(np.unique(y_true)) > 1:
+            if not (
+                len(np.unique(y_pred_proba)) < 2 and len(y_pred_proba) == len(y_true)
+            ):
+                try:
+                    roc_auc_val = roc_auc_score(y_true, y_pred_proba)
+                except ValueError:
+                    pass
+        metrics["roc_auc"] = roc_auc_val
+
+        metrics["classification_report_dict"] = classification_report(
+            y_true, y_pred_class, digits=4, zero_division=0, output_dict=True
+        )
+        metrics["confusion_matrix_list"] = confusion_matrix(
+            y_true, y_pred_class
+        ).tolist()
+        metrics["y_pred_test_classes_list"] = (
+            y_pred_class.tolist()
+            if isinstance(y_pred_class, np.ndarray)
+            else y_pred_class
+        )
+        metrics["y_prob_test_scores_list"] = (
+            y_pred_proba.tolist()
+            if y_pred_proba is not None and isinstance(y_pred_proba, np.ndarray)
+            else y_pred_proba
+        )
+    except Exception as e:
+        console.print(
+            f"[bold red]Error during evaluation for {model_name}: {e}[/bold red]"
+        )
+        for m_key in ["accuracy", "f1_score", "precision", "recall", "roc_auc"]:
+            metrics[m_key] = np.nan
+        metrics["classification_report_dict"] = {}
+        metrics["confusion_matrix_list"] = []
+    return metrics
+
+
+def plot_roc_curves_all(
+    experiment_results_dict: Dict[str, Dict[str, Any]],
+    y_true_labels: np.ndarray,
+    output_dir_path: str,
+    media_type: str,
+):
+    plt.figure(figsize=(12, 10))
+    plot_count = 0
+    for model_key, result_data in experiment_results_dict.items():
+        if (
+            "y_prob_test_scores_list" in result_data
+            and result_data["y_prob_test_scores_list"] is not None
+        ):
+            proba_scores = np.array(result_data["y_prob_test_scores_list"])
+            if len(np.unique(y_true_labels)) < 2 or (
+                proba_scores.ndim > 0
+                and len(np.unique(proba_scores)) < 2
+                and len(proba_scores) == len(y_true_labels)
+            ):
+                continue
+            try:
+                fpr, tpr, _ = roc_curve(y_true_labels, proba_scores)
+                roc_auc_value = result_data.get("roc_auc", auc(fpr, tpr))
+                if pd.notna(roc_auc_value):
+                    plt.plot(
+                        fpr,
+                        tpr,
+                        lw=1.8,
+                        label=f"{model_key} (AUC = {roc_auc_value:.4f})",
+                    )
+                    plot_count += 1
+            except ValueError as e:
+                console.print(
+                    f"[yellow]Could not plot ROC for {model_key} ({media_type}): {e}[/yellow]"
+                )
+
+    if plot_count > 0:
+        plt.plot([0, 1], [0, 1], color="grey", lw=1.5, linestyle="--")
+        plt.xlim([-0.01, 1.0])
+        plt.ylim([0.0, 1.01])
+        plt.xlabel("False Positive Rate", fontsize=13)
+        plt.ylabel("True Positive Rate", fontsize=13)
+        plt.title(
+            f"Meta-Learner & Ensemble ROC Curves ({media_type.capitalize()})",
+            fontsize=15,
+        )
+        plt.legend(loc="lower right", fontsize="small", frameon=True)
+        plt.grid(alpha=0.35, linestyle=":")
+        plt.tight_layout()
+        plot_path = os.path.join(
+            output_dir_path, f"all_meta_learners_roc_curves_{media_type}.png"
+        )
+        plt.savefig(plot_path, dpi=150)
+        console.print(
+            f"Combined ROC curves plot for {media_type} saved to [green]{plot_path}[/green]"
+        )
+    else:
+        console.print(f"[yellow]No valid ROC curves to plot for {media_type}.[/yellow]")
+    plt.close()
+
+
+def optimize_average_weights_simple_grid(
+    X_val_probs: np.ndarray,
+    y_val_true: np.ndarray,
+    num_base_models: int,
+    weight_options: Optional[List[float]] = None,
+) -> np.ndarray:
+    if weight_options is None:
+        weight_options = [0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0]
+    best_auc_val = -1.0
+    best_weights_val = np.ones(num_base_models)
+
+    max_combinations_exhaustive = 5**4
+    num_random_samples_if_large = 2000
+
+    if num_base_models <= 0:
+        console.print(
+            "[yellow]No base models to optimize weights for. Returning default weights.[/yellow]"
+        )
+        return best_weights_val
+
+    if (
+        num_base_models <= 4
+        and len(weight_options) ** num_base_models <= max_combinations_exhaustive
+    ):
+        weight_candidates = list(
+            itertools.product(weight_options, repeat=num_base_models)
+        )
+        console.print(
+            f"Optimizing average weights with exhaustive grid search ({len(weight_candidates)} trials)."
+        )
+    else:
+        console.print(
+            f"[yellow]Optimizing average weights with random sampling ({num_random_samples_if_large} trials due to {num_base_models} models).[/yellow]"
+        )
+        weight_candidates = [
+            np.array(np.random.choice(weight_options, num_base_models))
+            for _ in range(num_random_samples_if_large)
+        ]
+
+    with Progress(
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        BarColumn(),
+        TextColumn("{task.percentage:>3.1f}%"),
+        TimeElapsedColumn(),
+        MofNCompleteColumn(),
+    ) as progress:
+        task = progress.add_task("Weight Grid Search", total=len(weight_candidates))
+        for current_weights_tuple in weight_candidates:
+            current_weights = np.array(current_weights_tuple)
+            if np.sum(current_weights) == 0:
+                progress.update(task, advance=1)
+                continue
+
+            if X_val_probs.shape[0] == 0:
+                progress.update(task, advance=1)
+                continue
+            weighted_avg_probs_val_set = np.average(
+                X_val_probs, axis=1, weights=current_weights
+            )
+
+            current_auc_val = 0.0
+            if len(np.unique(y_val_true)) > 1 and not (
+                len(np.unique(weighted_avg_probs_val_set)) < 2
+                and len(weighted_avg_probs_val_set) == len(y_val_true)
+            ):
+                try:
+                    current_auc_val = roc_auc_score(
+                        y_val_true, weighted_avg_probs_val_set
+                    )
+                except ValueError:
+                    pass
+            if current_auc_val > best_auc_val:
+                best_auc_val, best_weights_val = current_auc_val, current_weights
+            progress.update(task, advance=1)
+
+    console.print(
+        f"Best weights from validation grid search: {best_weights_val.tolist()} with Val AUC: {best_auc_val:.4f}"
+    )
+    return best_weights_val
+
+
+# --- Main Experimentation Function ---
+def run_meta_learning_experiments(
+    meta_features_file: str,
+    output_dir_base: str,
+    api_artifacts_dir: str,
+    media_type: str,
+    optimizer_type: str,
+    n_optuna_trials_config: int,
+    provided_custom_weights: Optional[Dict[str, float]] = None,
+):
+    global OPTIMIZER_CHOICE, N_OPTUNA_TRIALS
+    OPTIMIZER_CHOICE = optimizer_type
+    N_OPTUNA_TRIALS = n_optuna_trials_config
+
+    if OPTIMIZER_CHOICE == "optuna" and not OPTIMIZER_AVAILABLE_OPTUNA:
+        console.print(
+            "[yellow]Optuna chosen but not installed. Falling back to GridSearchCV.[/yellow]"
+        )
+        OPTIMIZER_CHOICE = "gridsearch"
+
+    experiment_run_output_dir = os.path.join(
+        output_dir_base, f"experiments_{media_type}_{time.strftime('%Y%m%d_%H%M%S')}"
+    )
+    os.makedirs(experiment_run_output_dir, exist_ok=True)
+
+    # Main API artifacts directory (parent for media-specific subfolders)
+    os.makedirs(api_artifacts_dir, exist_ok=True)
+    # Media-type specific subdirectory within the main api_artifacts_dir
+    media_type_api_artifacts_subdir = os.path.join(api_artifacts_dir, media_type)
+    os.makedirs(media_type_api_artifacts_subdir, exist_ok=True)
+
+    console.rule(
+        f"[bold cyan]DeepSafe Meta-Learning: {media_type.upper()} (Optimizer: {OPTIMIZER_CHOICE})[/bold cyan]"
+    )
+    console.print(
+        Panel(
+            f"Meta-features: {meta_features_file}\n"
+            f"Experiment outputs: {os.path.abspath(experiment_run_output_dir)}\n"
+            f"API artifacts subfolder: {os.path.abspath(media_type_api_artifacts_subdir)}",
+            title="Paths",
+            border_style="dim blue",
+            expand=False,
+        )
+    )
+    all_experiment_results: Dict[str, Dict[str, Any]] = {}
+
+    console.rule("[bold]1. Data Loading and Preprocessing[/bold]")
+    try:
+        df_meta = pd.read_csv(meta_features_file)
+        console.print(
+            f"Loaded {media_type} meta-features from: [cyan]{meta_features_file}[/cyan], shape: {df_meta.shape}"
+        )
+    except Exception as e:
+        console.print(
+            f"[bold red]Fatal Error: Could not load meta-features file: {e}[/bold red]"
+        )
+        return
+
+    base_model_prob_features = sorted(
+        [col for col in df_meta.columns if col.endswith("_prob")]
+    )
+    if not base_model_prob_features:
+        console.print(
+            "[bold red]Fatal Error: No base model probability columns (ending with '_prob') found in CSV.[/bold red]"
+        )
+        return
+
+    console.print(
+        f"Identified [magenta]{len(base_model_prob_features)}[/magenta] base model probability features: {base_model_prob_features}"
+    )
+
+    temp_exp_feature_cols_path = os.path.join(
+        experiment_run_output_dir, f"experiment_feature_columns_{media_type}.json"
+    )
+    with open(temp_exp_feature_cols_path, "w") as f:
+        json.dump(base_model_prob_features, f, indent=2)
+
+    X_meta_all = df_meta[base_model_prob_features].copy()
+    y_meta_all = df_meta["ground_truth"]
+
+    cols_to_drop_all_nan = X_meta_all.columns[X_meta_all.isnull().all()].tolist()
+    if cols_to_drop_all_nan:
+        console.print(
+            f"[yellow]Warning: Dropping fully NaN columns: {cols_to_drop_all_nan}[/yellow]"
+        )
+        X_meta_all = X_meta_all.drop(columns=cols_to_drop_all_nan)
+        base_model_prob_features = [
+            col for col in base_model_prob_features if col not in cols_to_drop_all_nan
+        ]
+        if not base_model_prob_features:
+            console.print(
+                "[bold red]Fatal Error: All features became NaN after dropping some columns.[/bold red]"
+            )
+            return
+        with open(temp_exp_feature_cols_path, "w") as f:
+            json.dump(base_model_prob_features, f, indent=2)
+
+    X_meta_train_val, X_meta_test, y_meta_train_val, y_meta_test = train_test_split(
+        X_meta_all,
+        y_meta_all,
+        test_size=0.25,
+        random_state=42,
+        stratify=y_meta_all if len(np.unique(y_meta_all)) > 1 else None,
+    )
+    console.print(
+        f"Data split: Meta-Train/Val shape {X_meta_train_val.shape}, Meta-Test shape {X_meta_test.shape}"
+    )
+
+    ml_preprocessor = Pipeline(
+        [("imputer", SimpleImputer(strategy="median")), ("scaler", StandardScaler())]
+    )
+    X_meta_train_val_processed = ml_preprocessor.fit_transform(X_meta_train_val)
+    X_meta_test_processed = ml_preprocessor.transform(X_meta_test)
+
+    joblib.dump(
+        ml_preprocessor,
+        os.path.join(
+            experiment_run_output_dir, f"experiment_ml_preprocessor_{media_type}.joblib"
+        ),
+    )
+    console.print(
+        f"ML preprocessor for {media_type} (imputer + scaler) fitted and saved for this run."
+    )
+
+    imputer_for_simple_ensembles = ml_preprocessor.named_steps["imputer"]
+    X_meta_test_imputed_only_df = pd.DataFrame(
+        imputer_for_simple_ensembles.transform(X_meta_test), columns=X_meta_test.columns
+    )
+
+    console.rule("[bold]2. Defining ML Meta-Learners and Hyperparameter Spaces[/bold]")
+    models_and_param_spaces: Dict[str, Tuple[Any, Dict[str, Any]]] = {
+        "LogisticRegression": (
+            LogisticRegression(
+                solver="liblinear",
+                random_state=42,
+                class_weight="balanced",
+                max_iter=3000,
+            ),
+            {
+                "C": (
+                    (0.01, 1000.0, "loguniform")
+                    if OPTIMIZER_CHOICE == "optuna"
+                    else [0.01, 0.1, 1, 10, 100, 500]
+                )
+            },
+        ),
+        "RandomForest": (
+            RandomForestClassifier(random_state=42, class_weight="balanced"),
+            {
+                "n_estimators": (
+                    (100, 500, "int")
+                    if OPTIMIZER_CHOICE == "optuna"
+                    else [100, 200, 300, 400]
+                ),
+                "max_depth": (
+                    (5, 25, "int", True)
+                    if OPTIMIZER_CHOICE == "optuna"
+                    else [5, 10, 15, 20, None]
+                ),
+                "min_samples_split": (
+                    (2, 20, "int") if OPTIMIZER_CHOICE == "optuna" else [2, 5, 10, 15]
+                ),
+                "min_samples_leaf": (
+                    (1, 15, "int") if OPTIMIZER_CHOICE == "optuna" else [1, 5, 10, 15]
+                ),
+            },
+        ),
+        "GradientBoosting": (
+            GradientBoostingClassifier(random_state=42),
+            {
+                "n_estimators": (
+                    (100, 500, "int")
+                    if OPTIMIZER_CHOICE == "optuna"
+                    else [100, 200, 300, 400]
+                ),
+                "learning_rate": (
+                    (0.005, 0.2, "loguniform")
+                    if OPTIMIZER_CHOICE == "optuna"
+                    else [0.01, 0.05, 0.1, 0.15]
+                ),
+                "max_depth": (
+                    (3, 10, "int") if OPTIMIZER_CHOICE == "optuna" else [3, 5, 7, 9]
+                ),
+            },
+        ),
+        "SVC_Linear": (
+            SVC(
+                kernel="linear",
+                probability=True,
+                random_state=42,
+                class_weight="balanced",
+                max_iter=10000,
+            ),
+            {
+                "C": (
+                    (0.01, 100.0, "loguniform")
+                    if OPTIMIZER_CHOICE == "optuna"
+                    else [0.1, 1, 10, 100]
+                )
+            },
+        ),
+        "KNeighbors": (
+            KNeighborsClassifier(),
+            {
+                "n_neighbors": (
+                    (3, 25, "int", False, 2)
+                    if OPTIMIZER_CHOICE == "optuna"
+                    else [3, 5, 7, 11, 15, 19, 23]
+                ),
+                "weights": (
+                    (["uniform", "distance"], "categorical")
+                    if OPTIMIZER_CHOICE == "optuna"
+                    else ["uniform", "distance"]
+                ),
+            },
+        ),
+        "GaussianNB": (GaussianNB(), {}),
+    }
+    if XGBOOST_AVAILABLE and XGBClassifier:
+        models_and_param_spaces["XGBoost"] = (
+            XGBClassifier(random_state=42, eval_metric="auc"),
+            {
+                "n_estimators": (
+                    (100, 600, "int")
+                    if OPTIMIZER_CHOICE == "optuna"
+                    else [100, 200, 300, 400, 500]
+                ),
+                "learning_rate": (
+                    (0.005, 0.2, "loguniform")
+                    if OPTIMIZER_CHOICE == "optuna"
+                    else [0.01, 0.05, 0.1]
+                ),
+                "max_depth": (
+                    (3, 12, "int") if OPTIMIZER_CHOICE == "optuna" else [3, 5, 7, 9, 11]
+                ),
+                "scale_pos_weight": (
+                    (
+                        (np.sum(y_meta_train_val == 0) / np.sum(y_meta_train_val == 1))
+                        if np.sum(y_meta_train_val == 1) > 0
+                        else 1.0
+                    ),
+                ),
+            },
+        )
+    if LIGHTGBM_AVAILABLE and LGBMClassifier:
+        models_and_param_spaces["LightGBM"] = (
+            LGBMClassifier(
+                random_state=42, class_weight="balanced", metric="auc", verbosity=-1
+            ),
+            {
+                "n_estimators": (
+                    (100, 600, "int")
+                    if OPTIMIZER_CHOICE == "optuna"
+                    else [100, 200, 300, 400, 500]
+                ),
+                "learning_rate": (
+                    (0.005, 0.2, "loguniform")
+                    if OPTIMIZER_CHOICE == "optuna"
+                    else [0.01, 0.05, 0.1]
+                ),
+                "num_leaves": (
+                    (20, 150, "int")
+                    if OPTIMIZER_CHOICE == "optuna"
+                    else [31, 50, 70, 100, 130]
+                ),
+            },
+        )
+
+    console.rule(
+        f"[bold]3. Training and Evaluating ML-based Meta-Learners ({media_type.capitalize()} Stacking)[/bold]"
+    )
+    cv_strategy = StratifiedKFold(
+        n_splits=CV_FOLDS_DEFAULT, shuffle=True, random_state=42
+    )
+    trained_ml_model_objects: Dict[str, Any] = {}
+
+    for model_name_key, (
+        model_instance_template,
+        param_def,
+    ) in models_and_param_spaces.items():
+        console.rule(
+            f"[bold blue]Optimizing & Training {media_type.capitalize()} Meta-Learner: {model_name_key}[/bold blue]",
+            style="blue",
+        )
+        start_train_time = time.time()
+        best_estimator_for_model = None
+
+        if not param_def:
+            model_instance_template.fit(X_meta_train_val_processed, y_meta_train_val)
+            best_estimator_for_model = model_instance_template
+            console.print(
+                f"{model_name_key} fitted directly (no hyperparameters tuned)."
+            )
+        elif OPTIMIZER_CHOICE == "optuna" and optuna:
+
+            def optuna_objective(trial: optuna.Trial):
+                current_params = {}
+                for p_name, p_opts in param_def.items():
+                    if isinstance(p_opts, tuple) and len(p_opts) >= 2:
+                        suggestion_type_or_values = (
+                            p_opts[1]
+                            if p_name == "weights" and p_opts[1] == "categorical"
+                            else p_opts[2]
+                        )
+                        if suggestion_type_or_values == "loguniform":
+                            current_params[p_name] = trial.suggest_float(
+                                p_name, p_opts[0], p_opts[1], log=True
+                            )
+                        elif suggestion_type_or_values == "uniform":
+                            current_params[p_name] = trial.suggest_float(
+                                p_name, p_opts[0], p_opts[1]
+                            )
+                        elif suggestion_type_or_values == "int":
+                            low, high = p_opts[0], p_opts[1]
+                            can_be_none = p_opts[3] if len(p_opts) > 3 else False
+                            step = p_opts[4] if len(p_opts) > 4 else 1
+                            val = trial.suggest_int(p_name, low, high, step=step)
+                            if can_be_none and trial.suggest_categorical(
+                                f"{p_name}_use_none", [True, False]
+                            ):
+                                val = None
+                            current_params[p_name] = val
+                        elif suggestion_type_or_values == "categorical":
+                            current_params[p_name] = trial.suggest_categorical(
+                                p_name, p_opts[0]
+                            )
+                        elif len(p_opts) == 1 and not isinstance(p_opts[0], list):
+                            current_params[p_name] = p_opts[0]
+                        else:
+                            console.print(
+                                f"[red]Warning: Unknown Optuna parameter definition for {p_name}: {p_opts}[/red]"
+                            )
+                    else:
+                        if (
+                            p_name in model_instance_template.get_params()
+                            and not isinstance(p_opts, tuple)
+                        ):
+                            current_params[p_name] = p_opts
+
+                model_trial = model_instance_template.__class__(
+                    **model_instance_template.get_params()
+                )
+                valid_model_params = model_trial.get_params().keys()
+                filtered_current_params = {
+                    k: v for k, v in current_params.items() if k in valid_model_params
+                }
+                model_trial.set_params(**filtered_current_params)
+
+                scores = []
+                for train_idx, val_idx in cv_strategy.split(
+                    X_meta_train_val_processed, y_meta_train_val
+                ):
+                    X_fold_train, X_fold_val = (
+                        X_meta_train_val_processed[train_idx],
+                        X_meta_train_val_processed[val_idx],
+                    )
+                    y_fold_train, y_fold_val = (
+                        y_meta_train_val.iloc[train_idx],
+                        y_meta_train_val.iloc[val_idx],
+                    )
+                    model_trial.fit(X_fold_train, y_fold_train)
+                    if hasattr(model_trial, "predict_proba"):
+                        try:
+                            y_val_pred_proba = model_trial.predict_proba(X_fold_val)[
+                                :, 1
+                            ]
+                            if len(np.unique(y_fold_val)) < 2 or (
+                                len(np.unique(y_val_pred_proba)) < 2
+                                and len(y_val_pred_proba) == len(y_fold_val)
+                            ):
+                                scores.append(0.5)
+                            else:
+                                scores.append(
+                                    roc_auc_score(y_fold_val, y_val_pred_proba)
+                                )
+                        except Exception:
+                            scores.append(0.0)
+                    else:
+                        scores.append(
+                            f1_score(
+                                y_fold_val,
+                                model_trial.predict(X_fold_val),
+                                zero_division=0,
+                            )
+                        )
+                return np.mean(scores)
+
+            study = optuna.create_study(
+                direction="maximize", pruner=optuna.pruners.MedianPruner()
+            )
+            study.optimize(
+                optuna_objective,
+                n_trials=N_OPTUNA_TRIALS,
+                show_progress_bar=True,
+                gc_after_trial=True,
+            )
+
+            sklearn_best_params = {}
+            for p_name_orig_def, p_opts_def in param_def.items():
+                if p_name_orig_def in study.best_params:
+                    sklearn_best_params[p_name_orig_def] = study.best_params[
+                        p_name_orig_def
+                    ]
+                if len(p_opts_def) > 3 and p_opts_def[3] is True:
+                    if (
+                        study.best_params.get(f"{p_name_orig_def}_use_none", False)
+                        is True
+                    ):
+                        sklearn_best_params[p_name_orig_def] = None
+            console.print(
+                f"Best Optuna params for {model_name_key} ({media_type}): {sklearn_best_params}"
+            )
+            best_estimator_for_model = model_instance_template.__class__(
+                **model_instance_template.get_params()
+            )
+            best_estimator_for_model.set_params(**sklearn_best_params)
+            best_estimator_for_model.fit(X_meta_train_val_processed, y_meta_train_val)
+        else:
+            grid_search = GridSearchCV(
+                model_instance_template,
+                param_def,
+                cv=cv_strategy,
+                scoring="roc_auc",
+                n_jobs=-1,
+                verbose=0,
+            )
+            grid_search.fit(X_meta_train_val_processed, y_meta_train_val)
+            best_estimator_for_model = grid_search.best_estimator_
+            console.print(
+                f"Best GridSearchCV params for {model_name_key} ({media_type}): {grid_search.best_params_}"
+            )
+
+        joblib.dump(
+            best_estimator_for_model,
+            os.path.join(
+                experiment_run_output_dir,
+                f"{model_name_key}_meta_learner_{media_type}.joblib",
+            ),
+        )
+        trained_ml_model_objects[model_name_key] = best_estimator_for_model
+
+        y_test_pred_classes = best_estimator_for_model.predict(X_meta_test_processed)
+        y_test_pred_probas = (
+            best_estimator_for_model.predict_proba(X_meta_test_processed)[:, 1]
+            if hasattr(best_estimator_for_model, "predict_proba")
+            else None
+        )
+        metrics_results = evaluate_model_predictions(
+            y_meta_test.values, y_test_pred_classes, y_test_pred_probas, model_name_key
+        )
+        all_experiment_results[model_name_key] = metrics_results
+        train_time = time.time() - start_train_time
+        console.print(
+            f"[bold]{model_name_key} Test Set Perf. ({media_type}):[/bold] AUC: {metrics_results.get('roc_auc', np.nan):.4f}, F1: {metrics_results.get('f1_score', np.nan):.4f}, Acc: {metrics_results.get('accuracy', np.nan):.4f} (Train time: {train_time:.2f}s)"
+        )
+
+    console.rule(
+        f"[bold]4. Evaluating Simple Ensemble Baselines ({media_type.capitalize()} Meta-Test Set)[/bold]"
+    )
+    avg_probs_meta_test = X_meta_test_imputed_only_df.mean(axis=1).values
+    avg_preds_meta_test_classes = (
+        avg_probs_meta_test >= DEFAULT_THRESHOLD_FOR_SIMPLE_ENSEMBLES
+    ).astype(int)
+    all_experiment_results["Simple_Average_Prob"] = evaluate_model_predictions(
+        y_meta_test.values,
+        avg_preds_meta_test_classes,
+        avg_probs_meta_test,
+        "Simple_Average_Prob",
+    )
+    console.print(
+        f"[bold]Simple Average Prob Test ({media_type}):[/bold] AUC: {all_experiment_results['Simple_Average_Prob'].get('roc_auc', np.nan):.4f}, F1: {all_experiment_results['Simple_Average_Prob'].get('f1_score', np.nan):.4f}"
+    )
+
+    binarized_X_meta_test = (
+        X_meta_test_imputed_only_df.values >= DEFAULT_THRESHOLD_FOR_SIMPLE_ENSEMBLES
+    ).astype(int)
+    num_models_for_vote = X_meta_test_imputed_only_df.shape[1]
+    fake_votes_per_item_meta_test = binarized_X_meta_test.sum(axis=1)
+    maj_vote_preds_meta_test_classes = (
+        fake_votes_per_item_meta_test >= (num_models_for_vote / 2.0)
+    ).astype(int)
+    maj_vote_prob_scores_meta_test = (
+        fake_votes_per_item_meta_test / num_models_for_vote
+        if num_models_for_vote > 0
+        else np.full_like(fake_votes_per_item_meta_test, 0.5, dtype=float)
+    )
+    all_experiment_results["Simple_Majority_Vote"] = evaluate_model_predictions(
+        y_meta_test.values,
+        maj_vote_preds_meta_test_classes,
+        maj_vote_prob_scores_meta_test,
+        "Simple_Majority_Vote",
+    )
+    console.print(
+        f"[bold]Simple Majority Vote Test ({media_type}):[/bold] AUC: {all_experiment_results['Simple_Majority_Vote'].get('roc_auc', np.nan):.4f}, F1: {all_experiment_results['Simple_Majority_Vote'].get('f1_score', np.nan):.4f}"
+    )
+
+    if provided_custom_weights:
+        current_weights_values = [
+            provided_custom_weights.get(fc.replace("_prob", ""), 1.0)
+            for fc in base_model_prob_features
+        ]
+        current_weights_array = np.array(current_weights_values)
+
+        if (
+            len(current_weights_array) == X_meta_test_imputed_only_df.shape[1]
+            and np.sum(current_weights_array) > 0
+        ):
+            prov_weighted_avg_probs_meta_test = np.average(
+                X_meta_test_imputed_only_df.values,
+                axis=1,
+                weights=current_weights_array,
+            )
+            prov_weighted_avg_preds_meta_test_classes = (
+                prov_weighted_avg_probs_meta_test
+                >= DEFAULT_THRESHOLD_FOR_SIMPLE_ENSEMBLES
+            ).astype(int)
+            all_experiment_results["Provided_Weighted_Average"] = (
+                evaluate_model_predictions(
+                    y_meta_test.values,
+                    prov_weighted_avg_preds_meta_test_classes,
+                    prov_weighted_avg_probs_meta_test,
+                    "Provided_Weighted_Average",
+                )
+            )
+            console.print(
+                f"[bold]Provided Weighted Average Test ({media_type}):[/bold] AUC: {all_experiment_results['Provided_Weighted_Average'].get('roc_auc', np.nan):.4f}, F1: {all_experiment_results['Provided_Weighted_Average'].get('f1_score', np.nan):.4f}"
+            )
+        else:
+            console.print(
+                f"[yellow]Warning: Mismatch in provided_custom_weights keys vs. features for {media_type}, or sum of weights is zero. Skipping.[/yellow]"
+            )
+
+    X_train_val_imputed_for_opt_df = pd.DataFrame(
+        ml_preprocessor.named_steps["imputer"].transform(X_meta_train_val),
+        columns=base_model_prob_features,
+    )
+    stratify_opt_split = (
+        y_meta_train_val if len(np.unique(y_meta_train_val)) > 1 else None
+    )
+    X_opt_train_df, X_opt_val_df, y_opt_train_series, y_opt_val_series = (
+        train_test_split(
+            X_train_val_imputed_for_opt_df,
+            y_meta_train_val,
+            test_size=0.33,
+            random_state=123,
+            stratify=stratify_opt_split,
+        )
+    )
+    if X_opt_val_df.shape[0] > 10 and X_opt_val_df.shape[1] > 0:
+        console.print(
+            f"Optimizing weights for averaging ({media_type}) using a validation split of meta-train data..."
+        )
+        optimized_avg_weights = optimize_average_weights_simple_grid(
+            X_opt_val_df.values, y_opt_val_series.values, X_opt_val_df.shape[1]
+        )
+        opt_w_avg_probs_meta_test = np.average(
+            X_meta_test_imputed_only_df.values, axis=1, weights=optimized_avg_weights
+        )
+        opt_w_avg_preds_meta_test_classes = (
+            opt_w_avg_probs_meta_test >= DEFAULT_THRESHOLD_FOR_SIMPLE_ENSEMBLES
+        ).astype(int)
+        all_experiment_results["Optimized_Grid_Weighted_Average"] = (
+            evaluate_model_predictions(
+                y_meta_test.values,
+                opt_w_avg_preds_meta_test_classes,
+                opt_w_avg_probs_meta_test,
+                "Optimized_Grid_Weighted_Average",
+            )
+        )
+        console.print(
+            f"[bold]Optimized Grid Weighted Average Test ({media_type}):[/bold] AUC: {all_experiment_results['Optimized_Grid_Weighted_Average'].get('roc_auc', np.nan):.4f}, F1: {all_experiment_results['Optimized_Grid_Weighted_Average'].get('f1_score', np.nan):.4f}"
+        )
+
+        # Save optimized weights to media-type specific subdirectory with generic name
+        # (or keep media_type in name if preferred, but API loads generic name from subdir)
+        # opt_weights_api_path_generic = os.path.join(media_type_api_artifacts_subdir, "optimized_grid_average_weights.json")
+        # For now, keeping the original behavior of saving to main api_artifacts_dir with media_type in name
+        opt_weights_api_path_typed = os.path.join(
+            api_artifacts_dir, f"optimized_grid_average_weights_{media_type}.json"
+        )
+        with open(opt_weights_api_path_typed, "w") as f:
+            json.dump(
+                {
+                    feat: w
+                    for feat, w in zip(base_model_prob_features, optimized_avg_weights)
+                },
+                f,
+                indent=2,
+            )
+        console.print(
+            f"Optimized weights for {media_type} saved to API artifacts: [green]{opt_weights_api_path_typed}[/green]"
+        )
+    else:
+        console.print(
+            f"[yellow]Validation set for weight optimization ({media_type}) too small or no features. Skipping.[/yellow]"
+        )
+
+    console.rule(
+        f"[bold green]5. Overall Experiment Summary & Artifacts ({media_type.capitalize()})[/bold green]"
+    )
+    summary_table = Table(
+        title=f"Meta-Learner & Simple Ensemble Experiment Summary ({media_type.capitalize()} Meta-Test Set)"
+    )
+    summary_table.add_column(
+        "Method/Model", style="cyan", overflow="fold", max_width=35
+    )
+    summary_table.add_column("Test AUC", style="magenta")
+    summary_table.add_column("Test F1", style="green")
+    summary_table.add_column("Test Acc.", style="blue")
+    summary_table.add_column("Test Prec.", style="yellow")
+    summary_table.add_column("Test Recall", style="red")
+
+    sorted_results_list = sorted(
+        all_experiment_results.items(),
+        key=lambda item: (
+            item[1].get("roc_auc", -1) if pd.notna(item[1].get("roc_auc")) else -1
+        ),
+        reverse=True,
+    )
+    best_method_overall_name = "None"
+    best_method_overall_auc = -1.0
+    best_trainable_ml_model_for_api = None
+
+    for method_name_result, metrics_result in sorted_results_list:
+        summary_table.add_row(
+            method_name_result,
+            (
+                f"{metrics_result.get('roc_auc', 'N/A'):.4f}"
+                if pd.notna(metrics_result.get("roc_auc"))
+                else "N/A"
+            ),
+            f"{metrics_result.get('f1_score', 'N/A'):.4f}",
+            f"{metrics_result.get('accuracy', 'N/A'):.4f}",
+            f"{metrics_result.get('precision', 'N/A'):.4f}",
+            f"{metrics_result.get('recall', 'N/A'):.4f}",
+        )
+        current_auc_val_result = metrics_result.get("roc_auc", -1)
+        if (
+            pd.notna(current_auc_val_result)
+            and current_auc_val_result > best_method_overall_auc
+        ):
+            best_method_overall_auc = current_auc_val_result
+            best_method_overall_name = method_name_result
+            if method_name_result in trained_ml_model_objects:
+                best_trainable_ml_model_for_api = trained_ml_model_objects[
+                    method_name_result
+                ]
+
+    console.print(summary_table)
+    console.print(
+        f"\n[bold gold1]Best performing method overall for {media_type.upper()} (Test AUC): [white]{best_method_overall_name}[/white] (AUC: {best_method_overall_auc:.4f})[/bold gold1]"
+    )
+
+    results_json_path = os.path.join(
+        experiment_run_output_dir, f"all_experiments_metrics_summary_{media_type}.json"
+    )
+    with open(results_json_path, "w") as f:
+        json.dump(all_experiment_results, f, indent=2, cls=NpEncoder)
+    console.print(
+        f"All experiment metrics summaries for {media_type} saved to [green]{results_json_path}[/green]"
+    )
+
+    plot_roc_curves_all(
+        all_experiment_results,
+        y_meta_test.values,
+        experiment_run_output_dir,
+        media_type,
+    )
+
+    console.print(
+        f"\n[bold]Deployment Artifacts Preparation for {media_type.upper()} (in '{media_type_api_artifacts_subdir}'):[/bold]"
+    )
+
+    joblib.dump(
+        ml_preprocessor.named_steps["imputer"],
+        os.path.join(media_type_api_artifacts_subdir, "deepsafe_meta_imputer.joblib"),
+    )
+    joblib.dump(
+        ml_preprocessor.named_steps["scaler"],
+        os.path.join(media_type_api_artifacts_subdir, "deepsafe_meta_scaler.joblib"),
+    )
+
+    api_feature_cols_path = os.path.join(
+        media_type_api_artifacts_subdir, "deepsafe_meta_feature_columns.json"
+    )
+    if os.path.exists(temp_exp_feature_cols_path):
+        try:
+            with (
+                open(temp_exp_feature_cols_path, "r") as src_f,
+                open(api_feature_cols_path, "w") as dst_f,
+            ):
+                json.dump(json.load(src_f), dst_f, indent=2)
+            console.print(
+                f"Feature columns for {media_type} API saved to [green]{api_feature_cols_path}[/green]"
+            )
+        except Exception as e:
+            console.print(
+                f"[red]Error copying/saving feature columns file: {e}. Manual copy might be needed from {temp_exp_feature_cols_path} to {api_feature_cols_path}[/red]"
+            )
+    else:
+        console.print(
+            f"[yellow]Temporary feature columns file {temp_exp_feature_cols_path} not found. API artifact for feature columns may be missing for {media_type}.[/yellow]"
+        )
+
+    console.print(
+        f"Common imputer, scaler, and feature columns for {media_type} saved for API in '{media_type_api_artifacts_subdir}'."
+    )
+
+    if best_trainable_ml_model_for_api:
+        api_model_joblib_path = os.path.join(
+            media_type_api_artifacts_subdir, "deepsafe_meta_learner.joblib"
+        )
+        joblib.dump(best_trainable_ml_model_for_api, api_model_joblib_path)
+        console.print(
+            f"Best trainable ML meta-learner ([white]{best_method_overall_name}[/white]) for {media_type} saved as '{os.path.basename(api_model_joblib_path)}' in '{media_type_api_artifacts_subdir}'."
+        )
+        console.print(
+            f"The 4 artifacts in '{media_type_api_artifacts_subdir}' are ready for the API."
+        )
+    elif best_method_overall_name.startswith(("Simple", "Provided", "Optimized")):
+        console.print(
+            f"[yellow]The overall best method for {media_type} ([white]{best_method_overall_name}[/white]) is rule-based.[/yellow]"
+        )
+        console.print(
+            f"[yellow]To deploy a trainable ML model, choose the best one from this run and ensure its '.joblib' is saved as 'deepsafe_meta_learner.joblib' inside '{media_type_api_artifacts_subdir}'.[/yellow]"
+        )
+
+        opt_weights_main_dir_path = os.path.join(
+            api_artifacts_dir, f"optimized_grid_average_weights_{media_type}.json"
+        )
+        opt_weights_subdir_path_generic = os.path.join(
+            media_type_api_artifacts_subdir, "optimized_grid_average_weights.json"
+        )
+
+        if "Optimized_Grid_Weighted_Average" in best_method_overall_name:
+            if os.path.exists(opt_weights_main_dir_path):
+                console.print(
+                    f"  Optimized weights for this method are currently in '{opt_weights_main_dir_path}'. Consider standardizing its location if desired (e.g., to '{opt_weights_subdir_path_generic}')."
+                )
+            elif os.path.exists(
+                opt_weights_subdir_path_generic
+            ):  # If you adjust saving logic for weights too
+                console.print(
+                    f"  Optimized weights for this method are in '{opt_weights_subdir_path_generic}'."
+                )
+    else:
+        console.print(
+            f"[bold red]Error: Could not determine a best trainable model to save for {media_type}. Please review results.[/bold red]"
+        )
+
+    console.rule(
+        f"[bold green]Experimentation Suite for {media_type.upper()} Completed[/bold green]"
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Run Meta-Learning Experiments for DeepSafe Ensemble."
+    )
+    parser.add_argument(
+        "--media-type",
+        type=str,
+        choices=["image", "video", "audio"],
+        required=True,
+        help="Type of media for which the meta-learner is being trained (image, video, or audio).",
+    )
+    parser.add_argument(
+        "--meta-file",
+        type=str,
+        required=True,
+        help="Path to the media-specific meta-features CSV (e.g., ./meta_data/meta_features_image.csv)",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default=DEFAULT_EXPERIMENT_OUTPUT_DIR_BASE,
+        help=f"Base directory for saving all experiment-related outputs (default: {DEFAULT_EXPERIMENT_OUTPUT_DIR_BASE}).",
+    )
+    parser.add_argument(
+        "--api-artifacts-dir",
+        type=str,
+        default=DEFAULT_API_ARTIFACTS_DIR,
+        help=f"Directory to save final API-ready artifacts (default: {DEFAULT_API_ARTIFACTS_DIR})",
+    )
+    parser.add_argument(
+        "--optimizer",
+        type=str,
+        choices=["optuna", "gridsearch"],
+        default=OPTIMIZER_CHOICE_DEFAULT,
+        help=f"Hyperparameter optimizer (default: {OPTIMIZER_CHOICE_DEFAULT})",
+    )
+    parser.add_argument(
+        "--optuna-trials",
+        type=int,
+        default=N_OPTUNA_TRIALS_DEFAULT,
+        help=f"Number of Optuna trials (default: {N_OPTUNA_TRIALS_DEFAULT})",
+    )
+    parser.add_argument(
+        "--weights",
+        type=str,
+        default=None,
+        help='JSON string or path to JSON file for custom base model weights (for "Provided_Weighted_Average"). Keys should be base model names (e.g., "npr_deepfakedetection").',
+    )
+
+    args = parser.parse_args()
+
+    if OPTIMIZER_CHOICE_DEFAULT == "optuna" and not OPTIMIZER_AVAILABLE_OPTUNA:
+        console.print(
+            "[yellow]Default optimizer is Optuna, but it's not installed. GridSearchCV will be used if Optuna is chosen via CLI and not available.[/yellow]"
+        )
+    if not XGBOOST_AVAILABLE:
+        console.print(
+            "[yellow]XGBoost not installed. XGBoost experiments will be skipped if its block is reached.[/yellow]"
+        )
+    if not LIGHTGBM_AVAILABLE:
+        console.print(
+            "[yellow]LightGBM not installed. LightGBM experiments will be skipped if its block is reached.[/yellow]"
+        )
+
+    custom_weights_dict_main = None
+    if args.weights:
+        try:
+            if os.path.exists(args.weights):
+                with open(args.weights, "r") as f:
+                    custom_weights_dict_main = json.load(f)
+            else:
+                custom_weights_dict_main = json.loads(args.weights)
+            console.print(
+                f"Using provided custom base model weights: {custom_weights_dict_main}"
+            )
+        except Exception as e_weights:
+            console.print(
+                f"[bold red]Error parsing --weights argument: {e_weights}. Proceeding without them.[/bold red]"
+            )
+
+    run_meta_learning_experiments(
+        meta_features_file=args.meta_file,
+        output_dir_base=args.output_dir,
+        api_artifacts_dir=args.api_artifacts_dir,
+        media_type=args.media_type,
+        optimizer_type=args.optimizer,
+        n_optuna_trials_config=args.optuna_trials,
+        provided_custom_weights=custom_weights_dict_main,
+    )

image/aide/.gitignore

@@ -0,0 +1,5 @@
+weights/*.pth
+weights/*.pt
+weights/*/
+hf_cache/
+model_code/

image/aide/Dockerfile

@@ -0,0 +1,47 @@
+FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
+
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PYTHONUNBUFFERED=1
+
+WORKDIR /app
+
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends \
+    python3 python3-pip python3-dev \
+    git wget build-essential && \
+    rm -rf /var/lib/apt/lists/*
+
+RUN ln -sf /usr/bin/python3 /usr/bin/python
+
+RUN pip install --no-cache-dir --upgrade pip "setuptools>=68" wheel
+
+# Install PyTorch with CUDA 12.1 (replaces CPU-only arch-conditional install)
+RUN pip install --no-cache-dir \
+    torch==2.5.1 torchvision==0.20.1 \
+    --index-url https://download.pytorch.org/whl/cu121
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+RUN git clone https://github.com/shilinyan99/AIDE.git model_code && \
+    touch model_code/__init__.py && \
+    touch model_code/models/__init__.py && \
+    touch model_code/data/__init__.py
+
+RUN mkdir -p /app/weights /app/hf_cache
+COPY weights/ /app/weights/
+
+COPY app.py .
+
+ENV MODEL_PORT=5004
+ENV PRELOAD_MODEL=false
+ENV MODEL_TIMEOUT=600
+ENV AIDE_CHECKPOINT=GenImage_train.pth
+ENV PYTHONPATH=/app/model_code:$PYTHONPATH
+
+EXPOSE ${MODEL_PORT}
+
+RUN adduser --disabled-password --gecos '' appuser
+USER appuser
+
+CMD ["python", "app.py"]

image/aide/app.py

@@ -0,0 +1,433 @@
+import base64
+import gc
+import io
+import logging
+import os
+import platform
+import sys
+import threading
+import time
+from typing import Any, Dict, Optional
+
+import torch
+import uvicorn
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from PIL import Image, ImageFile
+from pydantic import BaseModel
+
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    handlers=[logging.StreamHandler(sys.stdout)],
+)
+logger = logging.getLogger(__name__)
+
+# ── Path setup ──────────────────────────────────────────────────────────────
+current_dir = os.path.dirname(os.path.abspath(__file__))
+model_code_dir = os.path.join(current_dir, "model_code")
+sys.path.insert(0, model_code_dir)
+sys.path.insert(0, os.path.join(model_code_dir, "models"))
+sys.path.insert(0, os.path.join(model_code_dir, "data"))
+
+# ── Compatibility shim ───────────────────────────────────────────────────────
+# AIDE's models/AIDE.py imports `clip` (openai-clip) at module level, but the
+# package is not used during inference — only open_clip is.  The openai-clip
+# package relies on pkg_resources which was removed in Python 3.13.  We inject
+# a lightweight stub so the import succeeds without installing the full package.
+import types as _types
+
+if "clip" not in sys.modules:
+    _clip_stub = _types.ModuleType("clip")
+    sys.modules["clip"] = _clip_stub
+
+# ── Config ──────────────────────────────────────────────────────────────────
+MODEL_NAME = "aide_detection"
+WEIGHTS_DIR = os.environ.get("WEIGHTS_DIR", os.path.join(current_dir, "weights"))
+HF_HOME = os.environ.get("HF_HOME", os.path.join(current_dir, "hf_cache"))
+os.environ["HF_HOME"] = HF_HOME
+# The checkpoint is self-contained (includes ConvNeXt weights), so we initialise
+# the architecture with no pretrained weights and load everything from the checkpoint.
+# Set CONVNEXT_PRETRAINED to a HuggingFace tag only if running without a checkpoint.
+CONVNEXT_PRETRAINED = os.environ.get("CONVNEXT_PRETRAINED", None)
+
+# Preferred checkpoint filename (GenImage trains on the most diverse generators)
+PREFERRED_CHECKPOINT = os.environ.get("AIDE_CHECKPOINT", "GenImage_train.pth")
+
+
+def _get_device():
+    """Select optimal device: MPS (Apple) > CUDA (NVIDIA) > CPU."""
+    override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower()
+    if override == "cpu":
+        return torch.device("cpu")
+    if override == "cuda" and torch.cuda.is_available():
+        return torch.device("cuda")
+    if (
+        override == "mps"
+        and hasattr(torch.backends, "mps")
+        and torch.backends.mps.is_available()
+    ):
+        return torch.device("mps")
+    if override:
+        pass  # Invalid override, fall through to auto-detect
+    if (
+        platform.system() == "Darwin"
+        and hasattr(torch.backends, "mps")
+        and torch.backends.mps.is_available()
+    ):
+        return torch.device("mps")
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    return torch.device("cpu")
+
+
+DEVICE = _get_device()
+if DEVICE.type == "cuda":
+    torch.backends.cudnn.benchmark = True
+    torch.set_float32_matmul_precision("high")
+if DEVICE.type == "cuda":
+    logger.info(
+        "Device: cuda (%s, %.1f GB VRAM)",
+        torch.cuda.get_device_name(0),
+        torch.cuda.get_device_properties(0).total_memory / 1024**3,
+    )
+else:
+    logger.warning(
+        "Device: %s (no CUDA available -- check nvidia-container-toolkit)",
+        DEVICE,
+    )
+PRELOAD_MODEL = os.environ.get("PRELOAD_MODEL", "false").lower() == "true"
+MODEL_TIMEOUT = int(os.environ.get("MODEL_TIMEOUT", "600"))
+
+# ── Globals ─────────────────────────────────────────────────────────────────
+model = None
+model_lock = threading.Lock()
+last_used_time = 0
+
+
+class ImageInput(BaseModel):
+    """Request body for /predict."""
+
+    image_data: str
+    threshold: Optional[float] = 0.5
+
+
+# ── Weight discovery ─────────────────────────────────────────────────────────
+
+
+def find_aide_checkpoint() -> Optional[str]:
+    """
+    Return path to the best AIDE checkpoint in WEIGHTS_DIR.
+
+    Priority order:
+      1. PREFERRED_CHECKPOINT filename (GenImage_train.pth by default)
+      2. Any other .pth file (largest wins)
+    """
+    if not os.path.exists(WEIGHTS_DIR):
+        logger.warning(f"Weights directory not found: {WEIGHTS_DIR}")
+        return None
+
+    # Try the preferred checkpoint first
+    preferred = os.path.join(WEIGHTS_DIR, PREFERRED_CHECKPOINT)
+    if os.path.exists(preferred):
+        logger.info(f"Using preferred checkpoint: {preferred}")
+        return preferred
+
+    # Fall back to the largest available checkpoint
+    candidates = [
+        os.path.join(WEIGHTS_DIR, f)
+        for f in os.listdir(WEIGHTS_DIR)
+        if f.endswith(".pth") or f.endswith(".pt")
+    ]
+    if not candidates:
+        logger.warning("No .pth checkpoint found in weights directory.")
+        return None
+    best = max(candidates, key=os.path.getsize)
+    logger.info(f"Using checkpoint: {best}")
+    return best
+
+
+# ── Preprocessing ────────────────────────────────────────────────────────────
+
+
+def preprocess_image(image_bytes: bytes) -> torch.Tensor:
+    """
+    Preprocess raw image bytes into AIDE's 5-view tensor.
+
+    Args:
+        image_bytes: Raw bytes of a JPEG/PNG/etc. image.
+
+    Returns:
+        Tensor of shape [1, 5, 3, 256, 256] on CPU.
+
+    Raises:
+        Exception: If bytes cannot be decoded or processed.
+    """
+    from data.dct import DCT_base_Rec_Module
+    from torchvision import transforms
+
+    pil_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+
+    # Ensure minimum 256x256 so DCT unfold has enough patches
+    w, h = pil_image.size
+    if w < 256 or h < 256:
+        pil_image = pil_image.resize((256, 256), Image.BICUBIC)
+
+    to_tensor = transforms.ToTensor()
+    image_tensor = to_tensor(pil_image)  # [3, H, W]
+
+    # DCT frequency decomposition → 4 patches [3, 32, 32] each
+    dct_module = DCT_base_Rec_Module()
+    x_minmin, x_maxmax, x_minmin1, x_maxmax1 = dct_module(image_tensor)
+
+    # Resize all views to 256×256 and normalise with ImageNet stats
+    transform = transforms.Compose(
+        [
+            transforms.Resize([256, 256]),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ]
+    )
+
+    x_0 = transform(image_tensor)
+    x_minmin = transform(x_minmin)
+    x_maxmax = transform(x_maxmax)
+    x_minmin1 = transform(x_minmin1)
+    x_maxmax1 = transform(x_maxmax1)
+
+    # Stack → [5, 3, 256, 256], unsqueeze batch → [1, 5, 3, 256, 256]
+    stacked = torch.stack([x_minmin, x_maxmax, x_minmin1, x_maxmax1, x_0], dim=0)
+    return stacked.unsqueeze(0).to(DEVICE)
+
+
+# ── Model loading ─────────────────────────────────────────────────────────────
+
+
+def load_model_internal():
+    """Load AIDE_Model onto CPU with the best available checkpoint."""
+    global model, last_used_time
+
+    with model_lock:
+        if model is not None:
+            last_used_time = time.time()
+            return
+
+        logger.info("Loading AIDE model...")
+        try:
+            import models.AIDE as AIDE_module
+
+            aide_model = AIDE_module.AIDE(
+                resnet_path=None,
+                convnext_path=CONVNEXT_PRETRAINED,
+            )
+            aide_model.to(DEVICE)
+
+            checkpoint_path = find_aide_checkpoint()
+            if checkpoint_path:
+                logger.info(f"Loading checkpoint: {checkpoint_path}")
+                ckpt = torch.load(checkpoint_path, map_location=DEVICE)
+                if isinstance(ckpt, dict):
+                    state_dict = ckpt.get("model") or ckpt.get("state_dict") or ckpt
+                else:
+                    state_dict = ckpt
+                # Strip DataParallel "module." prefix if present
+                cleaned = {
+                    k[7:] if k.startswith("module.") else k: v
+                    for k, v in state_dict.items()
+                }
+                missing, unexpected = aide_model.load_state_dict(cleaned, strict=False)
+                logger.info(
+                    f"Checkpoint loaded. Missing keys: {len(missing)}, "
+                    f"Unexpected keys: {len(unexpected)}"
+                )
+            else:
+                logger.warning(
+                    "No checkpoint found — model uses pretrained-only weights. "
+                    "Run download_weights.sh to fetch the fine-tuned checkpoint."
+                )
+
+            # Switch to inference mode (no gradient tracking, batch-norm uses running stats)
+            aide_model.train(mode=False)
+            model = aide_model
+            last_used_time = time.time()
+            logger.info("AIDE model ready.")
+
+        except Exception as exc:
+            logger.exception(f"Failed to load AIDE model: {exc}")
+            model = None
+            raise
+        finally:
+            gc.collect()
+
+
+def ensure_model_loaded():
+    """Load model on first request (lazy loading)."""
+    global last_used_time
+    if model is None:
+        load_model_internal()
+    else:
+        last_used_time = time.time()
+
+
+def unload_model_if_idle():
+    """Evict model from RAM after MODEL_TIMEOUT seconds of inactivity."""
+    global model
+    if model is None or PRELOAD_MODEL:
+        return
+    with model_lock:
+        if model is not None and (time.time() - last_used_time > MODEL_TIMEOUT):
+            logger.info("Unloading idle AIDE model to free RAM.")
+            del model
+            model = None
+            gc.collect()
+
+
+# ── FastAPI app ───────────────────────────────────────────────────────────────
+
+app = FastAPI(
+    title="AIDE Deepfake Detection Service",
+    description=(
+        "AI-generated image detection using AIDE (ICLR 2025) — "
+        "hybrid DCT frequency analysis + ConvNeXt-xxlarge semantic features."
+    ),
+    version="1.0.0",
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+@app.get("/")
+async def root():
+    """Root endpoint with service information."""
+    return {
+        "model_name": MODEL_NAME,
+        "description": "AIDE (ICLR 2025) AI-generated image detector",
+        "device": str(DEVICE),
+        "model_loaded": model is not None,
+    }
+
+
+def _gpu_health_info() -> dict:
+    """Return GPU metrics for the health endpoint."""
+    if torch.cuda.is_available() and DEVICE.type == "cuda":
+        return {
+            "gpu_name": torch.cuda.get_device_name(0),
+            "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2),
+            "vram_total_mb": round(
+                torch.cuda.get_device_properties(0).total_memory / 1024**2
+            ),
+        }
+    return {}
+
+
+@app.get("/health")
+async def health():
+    """Health check endpoint."""
+    return {
+        "status": "healthy",
+        "model_name": MODEL_NAME,
+        "device": str(DEVICE),
+        "model_loaded": model is not None,
+        **_gpu_health_info(),
+    }
+
+
+@app.post("/unload")
+async def unload_model_endpoint():
+    """Manually unload the model to free RAM."""
+    global model
+    if model is None:
+        return {"status": "not_loaded"}
+    del model
+    model = None
+    gc.collect()
+    return {"status": "success", "message": "Model unloaded."}
+
+
+@app.post("/predict")
+async def predict(image_input: ImageInput) -> Dict[str, Any]:
+    """
+    Predict whether the submitted image is AI-generated.
+
+    Args:
+        image_input: Base64-encoded image and optional classification threshold.
+
+    Returns:
+        Dict with model name, fake probability, binary prediction, class label, and inference time.
+    """
+    try:
+        ensure_model_loaded()
+        if model is None:
+            raise HTTPException(status_code=503, detail="Model not loaded.")
+
+        start = time.time()
+
+        try:
+            image_bytes = base64.b64decode(image_input.image_data)
+            input_tensor = preprocess_image(image_bytes)
+        except Exception as exc:
+            raise HTTPException(status_code=400, detail=f"Invalid image data: {exc}")
+
+        with torch.no_grad():
+            logits = model(input_tensor)  # [1, 2]
+            probs = torch.softmax(logits, dim=-1)  # [1, 2]
+            probability_fake = probs[0, 1].item()
+
+        prediction = 1 if probability_fake >= image_input.threshold else 0
+        class_label = "fake" if prediction == 1 else "real"
+        inference_time = time.time() - start
+
+        logger.info(
+            f"Prediction: {class_label} (prob={probability_fake:.4f}, {inference_time:.3f}s)"
+        )
+
+        if not PRELOAD_MODEL and MODEL_TIMEOUT > 0:
+            threading.Timer(MODEL_TIMEOUT + 5.0, unload_model_if_idle).start()
+
+        return {
+            "model": MODEL_NAME,
+            "probability": float(probability_fake),
+            "prediction": int(prediction),
+            "class": class_label,
+            "inference_time": float(inference_time),
+        }
+
+    except HTTPException:
+        raise
+    except Exception as exc:
+        logger.exception(f"Prediction error: {exc}")
+        raise HTTPException(status_code=500, detail=str(exc))
+
+
+@app.on_event("startup")
+async def startup_event():
+    """Startup handler — preloads model if PRELOAD_MODEL=true, else lazy-loads."""
+    if PRELOAD_MODEL:
+        logger.info("Preloading AIDE model at startup.")
+        try:
+            load_model_internal()
+        except Exception as exc:
+            logger.error(f"Preload failed: {exc}")
+    else:
+        logger.info("AIDE service ready — model loads on first request.")
+
+    if not PRELOAD_MODEL and MODEL_TIMEOUT > 0:
+
+        def _periodic_check():
+            unload_model_if_idle()
+            threading.Timer(MODEL_TIMEOUT / 2.0, _periodic_check).start()
+
+        threading.Timer(MODEL_TIMEOUT / 2.0, _periodic_check).start()
+
+
+if __name__ == "__main__":
+    port = int(os.environ.get("MODEL_PORT", 5004))
+    logger.info(f"Starting AIDE service on port {port}")
+    uvicorn.run(app, host="0.0.0.0", port=port)