Upload 6 files

app/api.py

@@ -0,0 +1,65 @@
+import os, io, base64
+import numpy as np
+from fastapi import FastAPI, UploadFile, File, Form
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from typing import Optional, Dict, Any
+from PIL import Image
+from matplotlib import cm
+
+BACKEND = os.getenv("DETECTOR_BACKEND", "wav2vec2").lower()
+try:
+    if BACKEND == "wav2vec2":
+        from .inference_wav2vec import Detector  # type: ignore
+    else:
+        from .inference import Detector  # type: ignore
+except Exception:
+    if BACKEND == "wav2vec2":
+        from app.inference_wav2vec import Detector  # type: ignore
+    else:
+        from app.inference import Detector  # type: ignore
+
+DEFAULT_WEIGHTS = "app/models/weights/wav2vec2_classifier.pth" if BACKEND=="wav2vec2" else "app/models/weights/cnn_melspec.pth"
+WEIGHTS = os.getenv("MODEL_WEIGHTS_PATH", DEFAULT_WEIGHTS)
+det = Detector(weights_path=WEIGHTS)
+
+app = FastAPI(title="Voice Guard API", version="1.1.0")
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # tighten in prod
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+class AnalyzeResponse(BaseModel):
+    human: float
+    ai: float
+    label: str
+    threshold: float
+    threshold_source: Optional[str] = None
+    backend: str
+    source_hint: str
+    replay_score: Optional[float] = None
+    decision: Optional[str] = None
+    decision_details: Optional[Dict[str, Any]] = None
+    heatmap_b64: str
+
+def heatmap_png_b64(cam: np.ndarray) -> str:
+    cam = np.clip(cam, 0.0, 1.0).astype(np.float32)
+    rgb = (cm.magma(cam)[..., :3] * 255).astype(np.uint8)
+    im = Image.fromarray(rgb)
+    buf = io.BytesIO(); im.save(buf, format="PNG")
+    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode("ascii")
+
+@app.post("/analyze", response_model=AnalyzeResponse)
+async def analyze(file: UploadFile = File(...), source_hint: str = Form("auto")):
+    raw = await file.read()
+    proba = det.predict_proba(raw, source_hint=source_hint)
+    cam = np.array(det.explain(raw, source_hint=source_hint)["cam"], dtype=np.float32)
+    return {
+        **proba,
+        "heatmap_b64": heatmap_png_b64(cam),
+    }
+
+@app.get("/health")
+def health(): return {"ok": True, "backend": BACKEND}

app/app.py

@@ -0,0 +1,138 @@
+import os
+import json
+import numpy as np
+import gradio as gr
+from dotenv import load_dotenv
+from matplotlib import cm
+ 
+load_dotenv()
+ 
+# -------------------------
+# 0) Env & defaults
+# -------------------------
+BACKEND = os.getenv("DETECTOR_BACKEND", "wav2vec2").strip().lower()  # "wav2vec2" or "cnn"
+ 
+DEFAULT_W2V_WEIGHTS = "app/models/weights/wav2vec2_classifier.pth"
+DEFAULT_CNN_WEIGHTS = "app/models/weights/cnn_melspec.pth"
+DEFAULT_WEIGHTS = DEFAULT_W2V_WEIGHTS if BACKEND == "wav2vec2" else DEFAULT_CNN_WEIGHTS
+MODEL_WEIGHTS_PATH = os.getenv("MODEL_WEIGHTS_PATH", DEFAULT_WEIGHTS).strip()
+ 
+ 
+# -------------------------
+# 1) Import your Detector
+# -------------------------
+def _import_detector(backend):
+    """
+    Import the correct Detector class depending on backend and package layout.
+    Works both when run as a module ('.inference_*') and as a script ('app.inference_*').
+    """
+    try:
+        if backend == "wav2vec2":
+            from .inference_wav2vec import Detector  # type: ignore
+        else:
+            from .inference import Detector  # type: ignore
+    except Exception:
+        if backend == "wav2vec2":
+            from app.inference_wav2vec import Detector  # type: ignore
+        else:
+            from app.inference import Detector  # type: ignore
+    return Detector
+ 
+ 
+try:
+    Detector = _import_detector(BACKEND)
+except Exception as e:
+    # Fallback dummy to keep the UI alive even if import fails,
+    # so you can see the error in the JSON panel.
+    class Detector:  # type: ignore
+        def __init__(self, *args, **kwargs):
+            self._err = f"Detector import failed: {e}"
+ 
+        def predict_proba(self, *args, **kwargs):
+            return {"error": self._err}
+ 
+        def explain(self, *args, **kwargs):
+            return {"cam": np.zeros((128, 128), dtype=np.float32).tolist()}
+ 
+ 
+# Single, shared detector (created lazily so startup is fast on Spaces)
+_DET = None
+def _get_detector():
+    global _DET
+    if _DET is None:
+        _DET = Detector(weights_path=MODEL_WEIGHTS_PATH)
+    return _DET
+ 
+ 
+# -------------------------
+# 2) Core functions
+# -------------------------
+def predict_and_explain(audio_path: str | None, source_hint: str):
+    """
+    audio_path: filepath from Gradio (since type='filepath')
+    source_hint: "Auto", "Microphone", "Upload"
+    """
+    source = (source_hint or "Auto").strip().lower()
+ 
+    if not audio_path or not os.path.exists(audio_path):
+        return {"error": "No audio received. Record or upload a 2–4s clip."}, None
+ 
+    det = _get_detector()
+ 
+    # Your Detector is expected to accept a file path and optional source hint
+    proba = det.predict_proba(audio_path, source_hint=source)
+    exp = det.explain(audio_path, source_hint=source)
+ 
+    # Explanation to heatmap (float [0,1] -> magma RGB uint8)
+    cam = np.array(exp.get("cam", []), dtype=np.float32)
+    if cam.ndim == 1:
+        # if model returned a 1D vector, tile to square-ish map
+        side = int(np.sqrt(cam.size))
+        side = max(side, 2)
+        cam = cam[: side * side].reshape(side, side)
+    cam = np.clip(cam, 0.0, 1.0)
+    cam_rgb = (cm.magma(cam)[..., :3] * 255).astype(np.uint8)
+ 
+    # Ensure proba is JSON-serializable
+    if not isinstance(proba, dict):
+        proba = {"result": proba}
+ 
+    return proba, cam_rgb
+ 
+ 
+def provenance(audio_path: str | None):
+    # Stub (you can wire a provenance model or checksum here)
+    return {"ok": True, "note": "Provenance check not wired in this app.py."}
+ 
+ 
+# -------------------------
+# 3) UI
+# -------------------------
+with gr.Blocks(title=f"AI Voice Detector · {BACKEND.upper()}") as demo:
+    gr.Markdown(f"# 🔎 AI Voice Detector — Backend: **{BACKEND.upper()}**")
+    gr.Markdown(
+        "Record or upload a short clip (~3s). Get probabilities, a label, and an explanation heatmap."
+    )
+ 
+    with gr.Row():
+        audio_in = gr.Audio(sources=["microphone", "upload"], type="filepath", label="Audio")
+        with gr.Column():
+            src = gr.Radio(choices=["Auto", "Microphone", "Upload"], value="Auto", label="Source")
+            btn_predict = gr.Button("Analyze", variant="primary")
+            btn_prov = gr.Button("Provenance Check (optional)")
+ 
+    with gr.Row():
+        json_out = gr.JSON(label="Prediction (probabilities + label)")
+        cam_out = gr.Image(label="Explanation Heatmap (saliency)")
+    prov_out = gr.JSON(label="Provenance Result (if available)")
+ 
+    btn_predict.click(predict_and_explain, inputs=[audio_in, src], outputs=[json_out, cam_out])
+    btn_prov.click(provenance, inputs=audio_in, outputs=prov_out)
+ 
+ 
+# -------------------------
+# 4) Launch (Spaces-friendly)
+# -------------------------
+if __name__ == "__main__":
+    # queue() keeps UI responsive under load; host/port are Spaces-safe and local-friendly
+    demo.queue().launch(server_name="0.0.0.0", server_port=7860)

app/elevenlabs_tools.py

@@ -0,0 +1,38 @@
+
+import os, time, hashlib, json, pathlib, random
+from typing import List, Optional
+from dotenv import load_dotenv
+import requests
+
+load_dotenv()
+ELEVEN_API_KEY = os.getenv("ELEVEN_API_KEY", "")
+ELEVEN_VOICE_ID = os.getenv("ELEVEN_VOICE_ID", "")
+BASE = "https://api.elevenlabs.io/v1"
+
+def _headers():
+    return {"xi-api-key": ELEVEN_API_KEY, "accept": "audio/mpeg", "Content-Type": "application/json"}
+
+def generate_tts_dataset(texts: List[str], voice_id: Optional[str]=None, out_dir: str="data/raw/ai", model_id: str="eleven_monolingual_v1"):
+    """Generate AI speech MP3s from ElevenLabs into out_dir. Convert to WAV (16k mono) for training."""
+    voice_id = voice_id or ELEVEN_VOICE_ID
+    assert ELEVEN_API_KEY, "Set ELEVEN_API_KEY in .env"
+    assert voice_id, "Provide ELEVEN_VOICE_ID in .env or pass voice_id"
+    os.makedirs(out_dir, exist_ok=True)
+    for i, txt in enumerate(texts):
+        payload = {"text": txt, "model_id": model_id, "voice_settings": {"stability": 0.4, "similarity_boost": 0.7}}
+        url = f"{BASE}/text-to-speech/{voice_id}"
+        r = requests.post(url, headers=_headers(), json=payload)
+        if r.status_code != 200:
+            print("TTS error", r.status_code, r.text[:200]); continue
+        mp3_path = os.path.join(out_dir, f"elab_{i:04d}.mp3")
+        with open(mp3_path, "wb") as f:
+            f.write(r.content)
+        print("saved", mp3_path)
+    print("Done. Convert MP3 to WAV (16kHz mono) before training.")
+
+def check_ai_speech(audio_bytes: bytes) -> dict:
+    """Stub: if your plan exposes classifier API, call it here; else returns unsupported."""
+    return {"supported": False, "prob_ai": None, "provider": "elevenlabs", "note": "Classifier not enabled in this template."}
+
+# if __name__ == "__main__":
+#     generate_tts_dataset()

app/inference_wav2vec.py

@@ -0,0 +1,214 @@
+import os, json
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from .models.wav2vec_detector import Wav2VecClassifier
+from .utils.audio import load_audio, pad_or_trim, TARGET_SR
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# ---------- Thresholds & biases ----------
+AI_THRESHOLD_DEFAULT = float(os.getenv("DETECTOR_AI_THRESHOLD", "0.60"))
+MIC_THRESHOLD        = float(os.getenv("DETECTOR_MIC_THRESHOLD", "0.68"))
+UPLOAD_THRESHOLD     = float(os.getenv("DETECTOR_UPLOAD_THRESHOLD", str(AI_THRESHOLD_DEFAULT)))
+AI_LOGIT_BIAS        = float(os.getenv("DETECTOR_AI_LOGIT_BIAS", "0.00"))  # add to AI logit globally
+
+# ---------- Decision rule ----------
+# 'threshold'  -> AI if ai_prob >= threshold
+# 'argmax'     -> AI if ai_prob > human_prob
+# 'hybrid'     -> threshold, but if replay_score >= T1 and ai_prob >= 0.50 -> AI
+DECISION_RULE       = os.getenv("DECISION_RULE", "threshold").lower()
+
+# ---------- Replay-attack heuristic ----------
+REPLAY_ENABLE       = os.getenv("REPLAY_ENABLE", "1") != "0"
+REPLAY_AI_BONUS     = float(os.getenv("REPLAY_AI_BONUS", "1.2"))
+REPLAY_FORCE_LABEL  = os.getenv("REPLAY_FORCE_LABEL", "0") == "1"
+REPLAY_T1           = float(os.getenv("REPLAY_T1", "0.35"))  # soft start
+REPLAY_T2           = float(os.getenv("REPLAY_T2", "0.55"))  # strong replay
+
+# ---------- DSP helpers ----------
+def peak_normalize(y: np.ndarray, peak: float = 0.95, eps: float = 1e-9) -> np.ndarray:
+    m = float(np.max(np.abs(y)) + eps)
+    return (y / m) * peak if m > 0 else y
+
+def rms_normalize(y: np.ndarray, target_rms: float = 0.03, eps: float = 1e-9) -> np.ndarray:
+    rms = float(np.sqrt(np.mean(y**2)) + eps)
+    g = target_rms / rms
+    return np.clip(y * g, -1.0, 1.0)
+
+def trim_silence(y: np.ndarray, sr: int, thresh_db: float = 40.0, min_ms: int = 30) -> np.ndarray:
+    if y.size == 0: return y
+    win = max(1, int(sr * 0.02))
+    pad = max(1, int(sr * (min_ms / 1000.0)))
+    energy = np.convolve(y ** 2, np.ones(win) / win, mode="same")
+    ref = np.max(energy) + 1e-12
+    mask = 10.0 * np.log10(energy / ref + 1e-12) > -thresh_db
+    if not np.any(mask): return y
+    idx = np.where(mask)[0]
+    start = max(0, int(idx[0] - pad))
+    end = min(len(y), int(idx[-1] + pad))
+    return y[start:end]
+
+def noise_gate(y, sr, gate_db=-42.0):
+    m = np.max(np.abs(y)) + 1e-9
+    thr = m * (10.0 ** (gate_db / 20.0))
+    y2 = y.copy()
+    y2[np.abs(y2) < thr] = 0.0
+    return y2
+
+def bandpass_fft(y: np.ndarray, sr: int, low=100.0, high=3800.0):
+    n = int(2 ** np.ceil(np.log2(len(y) + 1)))
+    Y = np.fft.rfft(y, n=n)
+    freqs = np.fft.rfftfreq(n, d=1.0/sr)
+    mask = (freqs >= low) & (freqs <= high)
+    Y_filtered = Y * mask
+    y_filt = np.fft.irfft(Y_filtered, n=n)[:len(y)]
+    return y_filt.astype(np.float32, copy=False)
+
+# ---------- Replay score ----------
+def replay_score(y: np.ndarray, sr: int) -> float:
+    if len(y) < sr:
+        y = np.pad(y, (0, sr - len(y)))
+    N = 4096
+    if len(y) < N:
+        y = np.pad(y, (0, N - len(y)))
+    w = np.hanning(N)
+    seg = y[:N] * w
+    X = np.abs(np.fft.rfft(seg)) + 1e-9
+
+    cep = np.fft.irfft(np.log(X))
+    qmin = max(1, int(0.0003 * sr))
+    qmax = min(len(cep) - 1, int(0.0040 * sr))
+    cwin = np.abs(cep[qmin:qmax])
+    c_peak = float(np.max(cwin)); c_mean = float(np.mean(cwin) + 1e-9)
+    cep_score = np.clip((c_peak - c_mean) / (c_peak + c_mean), 0.0, 1.0)
+
+    F = np.fft.rfftfreq(N, 1.0 / sr)
+    total = float(np.sum(X))
+    hf = float(np.sum(X[F >= 5000.0]))
+    hf_ratio = hf / (total + 1e-9)
+    hf_term = np.clip((0.25 - hf_ratio) / 0.25, 0.0, 1.0)
+
+    return float(np.clip(0.6 * cep_score + 0.4 * hf_term, 0.0, 1.0))
+
+# ---------- Detector ----------
+class Detector:
+    def __init__(self, weights_path: str, encoder: str | None = None, unfreeze_last: int = 0):
+        cfg = None
+        js = weights_path.replace(".pth", ".json")
+        if os.path.exists(js):
+            try:
+                with open(js, "r", encoding="utf-8") as f:
+                    cfg = json.load(f)
+            except Exception:
+                cfg = None
+        enc = encoder or (cfg.get("encoder") if cfg else "facebook/wav2vec2-base")
+        unf = unfreeze_last or int(cfg.get("unfreeze_last", 0)) if cfg else 0
+
+        self.model = Wav2VecClassifier(encoder=enc, unfreeze_last=unf).to(DEVICE)
+        if weights_path and os.path.exists(weights_path):
+            state = torch.load(weights_path, map_location=DEVICE)
+            self.model.load_state_dict(state, strict=False)
+        self.model.eval()
+
+    def _preprocess(self, y: np.ndarray, sr: int, source_hint: str | None):
+        y = trim_silence(y, sr, 40.0, 30)
+        y = bandpass_fft(y, sr, 100.0, 3800.0)
+        if source_hint and source_hint.lower().startswith("micro"):
+            y = noise_gate(y, sr, -42.0)
+            y = rms_normalize(y, 0.035)
+            y = peak_normalize(y, 0.95)
+        else:
+            y = rms_normalize(y, 0.03)
+            y = peak_normalize(y, 0.95)
+        y = pad_or_trim(y, duration_s=3.0, sr=sr)
+        return y
+
+    @torch.inference_mode()
+    def predict_proba(self, wav_bytes_or_path, source_hint: str | None = None):
+        y0, sr = load_audio(wav_bytes_or_path, target_sr=TARGET_SR)
+        rscore = replay_score(y0, sr) if REPLAY_ENABLE else 0.0
+
+        y = self._preprocess(y0, sr, source_hint)
+        x = torch.from_numpy(y).float().unsqueeze(0).to(DEVICE)
+
+        logits, _ = self.model(x)
+        logits = logits.clone()
+        logits[:, 1] += AI_LOGIT_BIAS
+
+        # Replay bonus on AI logit
+        if REPLAY_ENABLE and (source_hint and source_hint.lower().startswith("micro")) and (rscore >= REPLAY_T1):
+            ramp = np.clip((rscore - REPLAY_T1) / max(REPLAY_T2 - REPLAY_T1, 1e-6), 0.0, 1.0)
+            logits[:, 1] += REPLAY_AI_BONUS * ramp
+
+        probs = F.softmax(logits, dim=-1).cpu().numpy()[0]
+        p_h, p_ai = float(probs[0]), float(probs[1])
+
+        thr_source = "mic" if (source_hint and source_hint.lower().startswith("micro")) else "upload"
+        thr = MIC_THRESHOLD if thr_source == "mic" else UPLOAD_THRESHOLD
+
+        # Labels by different rules
+        label_thresh = "ai" if p_ai >= thr else "human"
+        label_argmax = "ai" if p_ai > p_h else "human"
+        label_hybrid = label_thresh
+        if REPLAY_ENABLE and rscore >= REPLAY_T1 and p_ai >= 0.50:
+            label_hybrid = "ai"
+        if REPLAY_ENABLE and rscore >= REPLAY_T2 and (source_hint and source_hint.lower().startswith("micro")):
+            if REPLAY_FORCE_LABEL or p_ai >= (thr - 0.05):
+                label_hybrid = "ai"
+
+        if DECISION_RULE == "argmax":
+            label = label_argmax
+            rule_used = "argmax"
+        elif DECISION_RULE == "hybrid":
+            label = label_hybrid
+            rule_used = "hybrid(threshold+replay)"
+        else:
+            label = label_thresh
+            rule_used = "threshold"
+
+        return {
+            "human": p_h,
+            "ai": p_ai,
+            "label": label,
+            "threshold": float(thr),
+            "threshold_source": thr_source,
+            "backend": "wav2vec2",
+            "source_hint": (source_hint or "auto"),
+            "replay_score": float(rscore),
+            "decision": rule_used,
+            "decision_details": {
+                "ai_prob": p_ai,
+                "human_prob": p_h,
+                "prob_margin": p_ai - p_h,
+                "ai_vs_threshold_margin": p_ai - thr,
+                "replay_score": rscore,
+                "mic_threshold": MIC_THRESHOLD,
+                "upload_threshold": UPLOAD_THRESHOLD,
+                "force_label_AI": bool(REPLAY_FORCE_LABEL and rscore >= REPLAY_T2),
+            },
+        }
+
+    def explain(self, wav_bytes_or_path, source_hint: str | None = None):
+        self.model.eval()
+        y0, sr = load_audio(wav_bytes_or_path, target_sr=TARGET_SR)
+        y = self._preprocess(y0, sr, source_hint)
+        x = torch.from_numpy(y).float().unsqueeze(0).to(DEVICE)
+        x.requires_grad_(True)
+        logits, feats = self.model(x)
+        logits[:, 1].sum().backward(retain_graph=True)
+        if feats.grad is None:
+            s = x.grad.detach().abs().squeeze(0)
+            s = s / (s.max() + 1e-6)
+            H = 64
+            step = max(1, s.numel() // 256)
+            s_small = s[::step][:256].cpu().numpy()
+            cam = np.tile(s_small[None, :], (H, 1))
+        else:
+            g = feats.grad.detach().abs().sum(dim=-1).squeeze(0)
+            g = g / (g.max() + 1e-6)
+            H = 64
+            cam = np.tile(g.cpu().numpy()[None, :], (H, 1))
+        cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-6)
+        return {"cam": cam.tolist(), "probs": None}

app/train.py

@@ -0,0 +1,315 @@
+import os
+import argparse
+import random
+from pathlib import Path
+from contextlib import nullcontext
+import importlib.util
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torch.backends.cudnn as cudnn
+from torch.utils.data import Dataset, DataLoader
+
+# ---------- Local imports ----------
+try:
+    from .models.cnn_melspec import TinyMelCNN
+    from .utils.audio import load_audio, pad_or_trim, logmel, TARGET_SR
+except ImportError:
+    from app.models.cnn_melspec import TinyMelCNN
+    from app.utils.audio import load_audio, pad_or_trim, logmel, TARGET_SR
+
+# ---------- Augmentations (robust across versions) ----------
+from audiomentations import (
+    Compose, AddGaussianNoise, TimeStretch, PitchShift, BandPassFilter
+)
+
+def make_gain(min_db, max_db, p):
+    """Handle both min_gain_in_db/max_gain_in_db and min_gain_db/max_gain_db."""
+    from audiomentations import Gain as _Gain
+    try:
+        return _Gain(min_gain_in_db=min_db, max_gain_in_db=max_db, p=p)
+    except TypeError:
+        return _Gain(min_gain_db=min_db, max_gain_db=max_db, p=p)
+
+def make_clipping(p=0.3):
+    """
+    Build ClippingDistortion across versions.
+    Newer:  min_percent/max_percent (0..20 typical)
+    Older:  min_percentile_threshold/max_percentile_threshold in [0..100]
+    Returns None if not available.
+    """
+    try:
+        from audiomentations import ClippingDistortion as _Clip
+    except Exception:
+        return None
+
+    # Try newer signature
+    for kwargs in (
+        dict(min_percent=0.0, max_percent=20.0, p=p),
+        dict(min_percent=5.0,  max_percent=30.0, p=p),
+    ):
+        try:
+            return _Clip(**kwargs)
+        except Exception:
+            pass
+
+    # Try older signature
+    for kwargs in (
+        dict(min_percentile_threshold=95, max_percentile_threshold=100, p=p),
+        dict(min_percentile_threshold=90, max_percentile_threshold=99,  p=p),
+    ):
+        try:
+            return _Clip(**kwargs)
+        except Exception:
+            pass
+
+    return None
+
+def have_fast_mp3():
+    return importlib.util.find_spec("fast_mp3_augment") is not None
+
+def make_mp3_compression(min_bitrate=48, max_bitrate=96, p=0.6):
+    """
+    Only enable Mp3Compression when the fast backend is present.
+    On Windows without the extra package this often breaks; we skip it.
+    """
+    if not have_fast_mp3():
+        return None
+    try:
+        from audiomentations import Mp3Compression as _Mp3
+        # Prefer the fast backend; if API lacks backend arg, constructor still works.
+        try:
+            return _Mp3(min_bitrate=min_bitrate, max_bitrate=max_bitrate, p=p, backend="fast_mp3_augment")
+        except TypeError:
+            return _Mp3(min_bitrate=min_bitrate, max_bitrate=max_bitrate, p=p)
+    except Exception:
+        return None
+
+# ---------- Repro ----------
+def set_seed(seed: int = 42):
+    random.seed(seed); np.random.seed(seed)
+    torch.manual_seed(seed); torch.cuda.manual_seed_all(seed)
+
+# ---------- Dataset ----------
+class FolderDataset(Dataset):
+    """
+    data_dir/
+      human/*.wav
+      ai/*.wav
+    """
+    def __init__(self, root: str, split: str = "train", val_ratio: float = 0.15,
+                 seed: int = 42, clip_seconds: float = 3.0):
+        self.root = Path(root)
+        self.clip_seconds = float(clip_seconds)
+
+        human = sorted((self.root / "human").glob("*.wav"))
+        ai = sorted((self.root / "ai").glob("*.wav"))
+        pairs = [(p, 0) for p in human] + [(p, 1) for p in ai]
+
+        rng = random.Random(seed)
+        rng.shuffle(pairs)
+
+        n_val = int(len(pairs) * val_ratio)
+        self.items = pairs[n_val:] if split == "train" else pairs[:n_val]
+        self.is_train = split == "train"
+
+        self._len_h = sum(1 for _, y in self.items if y == 0)
+        self._len_a = sum(1 for _, y in self.items if y == 1)
+
+        # Human: mild, natural perturbations
+        self.aug_human = Compose([
+            AddGaussianNoise(min_amplitude=0.001, max_amplitude=0.01, p=0.4),
+            TimeStretch(min_rate=0.96, max_rate=1.04, p=0.3),
+            PitchShift(min_semitones=-1, max_semitones=1, p=0.2),
+            make_gain(-4, 4, p=0.3),
+        ])
+
+        # AI: replay-aware chain (speaker/room/mic simulation)
+        ai_transforms = [
+            BandPassFilter(min_center_freq=200.0, max_center_freq=3500.0, p=0.5),
+            AddGaussianNoise(min_amplitude=0.001, max_amplitude=0.01, p=0.3),
+            TimeStretch(min_rate=0.95, max_rate=1.05, p=0.25),
+            make_gain(-6, 6, p=0.3),
+        ]
+        clip = make_clipping(p=0.3)
+        if clip is not None:
+            ai_transforms.insert(1, clip)
+        mp3 = make_mp3_compression()
+        if mp3 is not None:
+            ai_transforms.insert(0, mp3)
+
+        self.aug_ai = Compose(ai_transforms)
+
+    def __len__(self):
+        return len(self.items)
+
+    def __getitem__(self, idx: int):
+        path, label = self.items[idx]
+        y, sr = load_audio(str(path), TARGET_SR)
+        y = pad_or_trim(y, duration_s=self.clip_seconds, sr=sr)
+
+        if self.is_train:
+            if label == 1:
+                y = self.aug_ai(samples=y, sample_rate=sr)
+            else:
+                y = self.aug_human(samples=y, sample_rate=sr)
+
+        mel = logmel(y, sr)  # (n_mels, T)
+        x = torch.from_numpy(mel).unsqueeze(0)  # (1, n_mels, T)
+        y_t = torch.tensor(label, dtype=torch.long)
+        return x, y_t
+
+# ---------- Dataloaders ----------
+def make_dataloaders(args):
+    ds_tr = FolderDataset(args.data_dir, split="train", val_ratio=args.val_ratio,
+                          seed=args.seed, clip_seconds=args.clip_seconds)
+    ds_va = FolderDataset(args.data_dir, split="val", val_ratio=args.val_ratio,
+                          seed=args.seed, clip_seconds=args.clip_seconds)
+
+    # Windows is happier with workers=0; keep configurable
+    workers = args.workers if args.workers >= 0 else (0 if os.name == "nt" else max(1, (os.cpu_count() or 4)//2))
+    pin = (not args.cpu) and torch.cuda.is_available()
+
+    dl_tr = DataLoader(
+        ds_tr, batch_size=args.batch_size, shuffle=True,
+        num_workers=workers, pin_memory=pin,
+        persistent_workers=(workers > 0), drop_last=True,
+    )
+    dl_va = DataLoader(
+        ds_va, batch_size=max(1, args.batch_size // 2), shuffle=False,
+        num_workers=workers, pin_memory=pin,
+        persistent_workers=(workers > 0),
+    )
+    return ds_tr, ds_va, dl_tr, dl_va
+
+def class_weights_from_dataset(ds: FolderDataset, eps: float = 1e-6):
+    n_h, n_a = max(ds._len_h, eps), max(ds._len_a, eps)
+    w_h = (n_h + n_a) / (2 * n_h)
+    w_a = (n_h + n_a) / (2 * n_a)
+    return torch.tensor([w_h, w_a], dtype=torch.float32)
+
+# ---------- Training / Eval ----------
+def train_one_epoch(model, dl, device, opt, scaler, autocast_ctx, loss_fn, grad_accum=1):
+    model.train()
+    total_loss = 0.0
+    correct = 0
+    seen = 0
+    opt.zero_grad(set_to_none=True)
+
+    for step, (x, y) in enumerate(dl):
+        x = x.to(device, non_blocking=True)
+        y = y.to(device, non_blocking=True)
+
+        with autocast_ctx:
+            logits = model(x)
+            loss = loss_fn(logits, y)
+
+        loss = loss / grad_accum
+        if getattr(scaler, "is_enabled", lambda: False)():
+            scaler.scale(loss).backward()
+        else:
+            loss.backward()
+
+        if (step + 1) % grad_accum == 0:
+            if getattr(scaler, "is_enabled", lambda: False)():
+                scaler.step(opt)
+                scaler.update()
+            else:
+                opt.step()
+            opt.zero_grad(set_to_none=True)
+
+        total_loss += float(loss) * x.size(0) * grad_accum
+        correct += int((logits.argmax(1) == y).sum().item())
+        seen += x.size(0)
+
+    return total_loss / max(seen, 1), correct / max(seen, 1)
+
+@torch.no_grad()
+def evaluate(model, dl, device, loss_fn):
+    model.eval()
+    total_loss = 0.0
+    correct = 0
+    seen = 0
+    for x, y in dl:
+        x = x.to(device, non_blocking=True)
+        y = y.to(device, non_blocking=True)
+        logits = model(x)
+        loss = loss_fn(logits, y)
+        total_loss += float(loss) * x.size(0)
+        correct += int((logits.argmax(1) == y).sum().item())
+        seen += x.size(0)
+    return total_loss / max(seen, 1), correct / max(seen, 1)
+
+def main(args):
+    set_seed(args.seed)
+    device = "cuda" if (torch.cuda.is_available() and not args.cpu) else "cpu"
+    cudnn.benchmark = True
+
+    ds_tr, ds_va, dl_tr, dl_va = make_dataloaders(args)
+    print(f"Train items: {len(ds_tr)} (human={ds_tr._len_h}, ai={ds_tr._len_a})")
+    print(f"Val   items: {len(ds_va)}")
+
+    model = TinyMelCNN().to(device)
+
+    weights = class_weights_from_dataset(ds_tr).to(device)
+    loss_fn = torch.nn.CrossEntropyLoss(weight=weights)
+
+    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4)
+
+    # AMP (use new torch.amp if available, else fallback)
+    try:
+        from torch.amp import GradScaler, autocast as amp_autocast
+        scaler = GradScaler("cuda", enabled=(device == "cuda" and args.amp))
+        autocast_ctx = amp_autocast("cuda") if (device == "cuda" and args.amp) else nullcontext()
+    except Exception:
+        from torch.cuda.amp import GradScaler, autocast as amp_autocast  # deprecated but works
+        scaler = GradScaler(enabled=(device == "cuda" and args.amp))
+        autocast_ctx = amp_autocast() if (device == "cuda" and args.amp) else nullcontext()
+
+    best_va = -1.0
+    patience_counter = 0
+    Path(args.out).parent.mkdir(parents=True, exist_ok=True)
+
+    for epoch in range(args.epochs):
+        tr_loss, tr_acc = train_one_epoch(
+            model, dl_tr, device, opt, scaler, autocast_ctx, loss_fn,
+            grad_accum=args.grad_accum
+        )
+        va_loss, va_acc = evaluate(model, dl_va, device, loss_fn)
+
+        print(f"epoch {epoch+1:02d}/{args.epochs} | train {tr_loss:.3f}/{tr_acc:.3f} | val {va_loss:.3f}/{va_acc:.3f}")
+
+        # Save "last" every epoch
+        torch.save(model.state_dict(), args.out.replace(".pth", ".last.pth"))
+
+        if va_acc > best_va + 1e-4:
+            best_va = va_acc
+            torch.save(model.state_dict(), args.out)
+            patience_counter = 0
+            print(f"✅ Saved best to {args.out} (val_acc={best_va:.3f})")
+        else:
+            patience_counter += 1
+            if args.early_stop > 0 and patience_counter >= args.early_stop:
+                print(f"⏹️ Early stopping at epoch {epoch+1} (best val_acc={best_va:.3f})")
+                break
+
+    print("Done.")
+
+if __name__ == "__main__":
+    p = argparse.ArgumentParser(description="Train AI Voice Detector (replay-aware, version-robust, no fast_mp3 required)")
+    p.add_argument("--data_dir", type=str, required=True, help="Folder with subfolders human/ and ai/")
+    p.add_argument("--out", type=str, default="app/models/weights/cnn_melspec.pth")
+    p.add_argument("--epochs", type=int, default=10)
+    p.add_argument("--batch_size", type=int, default=32)
+    p.add_argument("--grad_accum", type=int, default=2)
+    p.add_argument("--lr", type=float, default=1e-3)
+    p.add_argument("--val_ratio", type=float, default=0.15)
+    p.add_argument("--clip_seconds", type=float, default=3.0)
+    p.add_argument("--workers", type=int, default=-1)   # try --workers 0 on Windows if you see issues
+    p.add_argument("--amp", action="store_true", default=True)
+    p.add_argument("--cpu", action="store_true")
+    p.add_argument("--early_stop", type=int, default=0)
+    p.add_argument("--seed", type=int, default=42)
+    args = p.parse_args()
+    main(args)

app/train_wav2vec.py

@@ -0,0 +1,207 @@
+import os, json, argparse, random
+from pathlib import Path
+from contextlib import nullcontext
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torch.backends.cudnn as cudnn
+from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler
+
+try:
+    from .models.wav2vec_detector import Wav2VecClassifier
+    from .utils.audio import load_audio, pad_or_trim, TARGET_SR
+except ImportError:
+    from app.models.wav2vec_detector import Wav2VecClassifier
+    from app.utils.audio import load_audio, pad_or_trim, TARGET_SR
+
+from audiomentations import Compose, AddGaussianNoise, BandPassFilter
+def make_gain(min_db, max_db, p):
+    from audiomentations import Gain as _Gain
+    try: return _Gain(min_gain_in_db=min_db, max_gain_in_db=max_db, p=p)
+    except TypeError: return _Gain(min_gain_db=min_db, max_gain_db=max_db, p=p)
+
+def set_seed(s=42):
+    random.seed(s); np.random.seed(s); torch.manual_seed(s); torch.cuda.manual_seed_all(s)
+
+def peak_normalize(y: np.ndarray, peak: float = 0.95, eps: float = 1e-9) -> np.ndarray:
+    m = float(np.max(np.abs(y)) + eps)
+    return (y / m) * peak if m > 0 else y
+
+def trim_silence(y: np.ndarray, sr: int, thresh_db: float = 40.0, min_ms: int = 30) -> np.ndarray:
+    if y.size == 0: return y
+    win = max(1, int(sr * 0.02)); pad = max(1, int(sr * (min_ms / 1000.0)))
+    energy = np.convolve(y**2, np.ones(win)/win, mode="same")
+    ref = np.max(energy) + 1e-12
+    mask = 10*np.log10(energy/ref + 1e-12) > -thresh_db
+    if not np.any(mask): return y
+    idx = np.where(mask)[0]; start=max(0,int(idx[0]-pad)); end=min(len(y),int(idx[-1]+pad))
+    return y[start:end]
+
+class WavDataset(Dataset):
+    """data_dir/{human,ai}/*.wav"""
+    def __init__(self, root, split="train", val_ratio=0.15, seed=42, clip_seconds=3.0):
+        self.root = Path(root); self.clip = float(clip_seconds)
+        human = sorted((self.root/"human").glob("*.wav"))
+        ai    = sorted((self.root/"ai").glob("*.wav"))
+        items = [(p,0) for p in human] + [(p,1) for p in ai]
+        rng = random.Random(seed); rng.shuffle(items)
+        n_val = int(len(items)*val_ratio)
+        self.items = items[n_val:] if split=="train" else items[:n_val]
+        self.is_train = split=="train"
+        self.nh = sum(1 for _,y in self.items if y==0)
+        self.na = sum(1 for _,y in self.items if y==1)
+        self.aug_h = Compose([AddGaussianNoise(0.001,0.01,p=0.3), make_gain(-4,4,p=0.3)])
+        self.aug_a = Compose([BandPassFilter(200.0,3500.0,p=0.5), AddGaussianNoise(0.001,0.01,p=0.3), make_gain(-6,6,p=0.3)])
+
+    def __len__(self): return len(self.items)
+    def __getitem__(self, idx):
+        path, label = self.items[idx]
+        y, sr = load_audio(str(path), TARGET_SR)
+        y = trim_silence(y, sr, 40.0, 30)
+        y = peak_normalize(y, 0.95)
+        y = pad_or_trim(y, duration_s=self.clip, sr=sr)
+        if self.is_train:
+            y = (self.aug_a if label==1 else self.aug_h)(samples=y, sample_rate=sr)
+        return torch.from_numpy(y).float(), torch.tensor(label, dtype=torch.long)
+
+def make_loaders(args):
+    ds_tr = WavDataset(args.data_dir, "train", args.val_ratio, args.seed, args.clip_seconds)
+    ds_va = WavDataset(args.data_dir, "val",   args.val_ratio, args.seed, args.clip_seconds)
+
+    # Weighted sampler to balance classes
+    labels = [y for _, y in ds_tr.items]
+    n0 = max(1, labels.count(0)); n1 = max(1, labels.count(1))
+    w0 = (n0 + n1) / (2 * n0); w1 = (n0 + n1) / (2 * n1)
+    sample_weights = [w0 if y == 0 else w1 for y in labels]
+    sampler = WeightedRandomSampler(sample_weights, num_samples=len(labels), replacement=True)
+
+    workers = args.workers if args.workers >= 0 else (0 if os.name=="nt" else max(1,(os.cpu_count() or 4)//2))
+    pin = (not args.cpu) and torch.cuda.is_available()
+    dl_tr = DataLoader(ds_tr, batch_size=args.batch_size, sampler=sampler,
+                       num_workers=workers, pin_memory=pin, drop_last=True)
+    dl_va = DataLoader(ds_va, batch_size=max(1,args.batch_size//2), shuffle=False,
+                       num_workers=workers, pin_memory=pin)
+    return ds_tr, ds_va, dl_tr, dl_va
+
+class FocalLoss(torch.nn.Module):
+    def __init__(self, alpha=None, gamma=1.5):
+        super().__init__()
+        self.alpha = alpha
+        self.gamma = gamma
+        self.ce = torch.nn.CrossEntropyLoss(weight=alpha)
+    def forward(self, logits, target):
+        ce = self.ce(logits, target)
+        with torch.no_grad():
+            pt = torch.exp(-ce)
+        return ((1 - pt) ** self.gamma) * ce
+
+def train_one_epoch(model, dl, device, opt, scaler, autocast_ctx, loss_fn, grad_accum=1):
+    model.train(); total=0.0; correct=0; seen=0
+    opt.zero_grad(set_to_none=True)
+    for step,(x,y) in enumerate(dl):
+        x=x.to(device,non_blocking=True); y=y.to(device,non_blocking=True)
+        with autocast_ctx:
+            logits,_=model(x); loss=loss_fn(logits,y)
+        loss=loss/grad_accum
+        if getattr(scaler,"is_enabled",lambda:False)(): scaler.scale(loss).backward()
+        else: loss.backward()
+        if (step+1)%grad_accum==0:
+            if getattr(scaler,"is_enabled",lambda:False)():
+                scaler.step(opt); scaler.update()
+            else:
+                opt.step()
+            opt.zero_grad(set_to_none=True)
+        total += float(loss) * x.size(0) * grad_accum
+        correct += int((logits.argmax(1)==y).sum().item()); seen += x.size(0)
+    return total/max(seen,1), correct/max(seen,1)
+
+@torch.no_grad()
+def evaluate(model, dl, device, loss_fn):
+    model.eval(); total=0.0; correct=0; seen=0
+    for x,y in dl:
+        x=x.to(device,non_blocking=True); y=y.to(device,non_blocking=True)
+        logits,_=model(x); loss=loss_fn(logits,y)
+        total += float(loss) * x.size(0); correct += int((logits.argmax(1)==y).sum().item()); seen += x.size(0)
+    return total/max(seen,1), correct/max(seen,1)
+
+def main(args):
+    set_seed(args.seed)
+    device = "cuda" if (torch.cuda.is_available() and not args.cpu) else "cpu"
+    cudnn.benchmark = True
+
+    ds_tr, ds_va, dl_tr, dl_va = make_loaders(args)
+    print(f"Train items: {len(ds_tr)} (human={ds_tr.nh}, ai={ds_tr.na})")
+    print(f"Val   items: {len(ds_va)}")
+
+    model = Wav2VecClassifier(
+        encoder=args.encoder,
+        unfreeze_last=args.unfreeze_last,
+        dropout=args.dropout,
+        hidden=args.hidden
+    ).to(device)
+
+    # Focal loss with class weights
+    nh, na = ds_tr.nh, ds_tr.na
+    w = torch.tensor([(nh+na)/(2*nh+1e-6), (nh+na)/(2*na+1e-6)], dtype=torch.float32).to(device)
+    loss_fn = FocalLoss(alpha=w, gamma=1.5)
+
+    head_params = list(model.head.parameters())
+    enc_params  = [p for p in model.encoder.parameters() if p.requires_grad]
+    param_groups = [{"params": head_params, "lr": args.lr_head}]
+    if enc_params:
+        param_groups.append({"params": enc_params, "lr": args.lr_encoder})
+    opt = torch.optim.AdamW(param_groups, weight_decay=1e-4)
+
+    try:
+        from torch.amp import GradScaler, autocast as amp_autocast
+        scaler = GradScaler("cuda", enabled=(device=="cuda" and args.amp))
+        autocast_ctx = amp_autocast("cuda") if (device=="cuda" and args.amp) else nullcontext()
+    except Exception:
+        from torch.cuda.amp import GradScaler, autocast as amp_autocast
+        scaler = GradScaler(enabled=(device=="cuda" and args.amp))
+        autocast_ctx = amp_autocast() if (device=="cuda" and args.amp) else nullcontext()
+
+    best=-1.0; patience=0
+    Path(args.out).parent.mkdir(parents=True, exist_ok=True)
+    with open(args.out.replace(".pth",".json"), "w", encoding="utf-8") as f:
+        json.dump({"encoder": args.encoder, "unfreeze_last": args.unfreeze_last}, f)
+
+    for epoch in range(args.epochs):
+        tr_loss, tr_acc = train_one_epoch(model, dl_tr, device, opt, scaler, autocast_ctx, loss_fn, args.grad_accum)
+        va_loss, va_acc = evaluate(model, dl_va, device, loss_fn)
+        print(f"epoch {epoch+1:02d}/{args.epochs} | train {tr_loss:.3f}/{tr_acc:.3f} | val {va_loss:.3f}/{va_acc:.3f}")
+        torch.save(model.state_dict(), args.out.replace(".pth",".last.pth"))
+        if va_acc > best + 1e-4:
+            best = va_acc; patience=0
+            torch.save(model.state_dict(), args.out)
+            print(f"✅ Saved best to {args.out} (val_acc={best:.3f})")
+        else:
+            patience += 1
+            if args.early_stop>0 and patience>=args.early_stop:
+                print(f"⏹️ Early stopping at epoch {epoch+1} (best={best:.3f})")
+                break
+    print("Done.")
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser(description="Train Wav2Vec2-based AI Voice Detector (balanced)")
+    ap.add_argument("--data_dir", required=True, help="Folder with human/ and ai/ WAVs")
+    ap.add_argument("--out", default="app/models/weights/wav2vec2_classifier.pth")
+    ap.add_argument("--encoder", default="facebook/wav2vec2-base")
+    ap.add_argument("--unfreeze_last", type=int, default=0)
+    ap.add_argument("--epochs", type=int, default=8)
+    ap.add_argument("--batch_size", type=int, default=16)
+    ap.add_argument("--grad_accum", type=int, default=2)
+    ap.add_argument("--lr_head", type=float, default=1e-3)
+    ap.add_argument("--lr_encoder", type=float, default=1e-5)
+    ap.add_argument("--val_ratio", type=float, default=0.15)
+    ap.add_argument("--clip_seconds", type=float, default=3.0)
+    ap.add_argument("--workers", type=int, default=-1)
+    ap.add_argument("--amp", action="store_true", default=True)
+    ap.add_argument("--cpu", action="store_true")
+    ap.add_argument("--dropout", type=float, default=0.2)
+    ap.add_argument("--hidden", type=int, default=256)
+    ap.add_argument("--early_stop", type=int, default=0)
+    ap.add_argument("--seed", type=int, default=42)
+    args = ap.parse_args()
+    main(args)