Upload app.py with huggingface_hub

app.py

@@ -1,9 +1,9 @@
 """Kaeva Verify V10 — Ensemble Deepfake Detection API.
 
 Endpoints:
-  POST /image   — V10 ensemble deepfake image detection (platform-aware)
-  POST /audio   — Wav2Vec2 V10 deepfake audio detection (3-class: real/tts/vc)
-  POST /video   — Video frame + audio deepfake analysis
+  POST /image   — V10 ensemble image detection (platform-aware)
+  POST /audio   — Wav2Vec2-based deepfake audio detection  
+  POST /video   — Video frame + audio analysis
   GET  /health  — Health check
 """
 
@@ -11,428 +11,367 @@ import io, os, traceback, tempfile, subprocess, time, json
 import numpy as np
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 from fastapi import FastAPI, UploadFile, File, HTTPException, Query
 from fastapi.middleware.cors import CORSMiddleware
 from PIL import Image
-from torchvision.models import efficientnet_b0, efficientnet_b3, resnet18
+from torchvision import transforms
+from torchvision.models import efficientnet_b0, efficientnet_b3
+from transformers import Wav2Vec2Model, Wav2Vec2FeatureExtractor
+import librosa
 
-app = FastAPI(title="Kaeva Verify V10", version="2.0.0")
+app = FastAPI(title="Kaeva Verify V10", version="10.0.0")
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
 
-# ──────────────────────────────────────────────
-# Ensemble configs per platform
-# ──────────────────────────────────────────────
+DEVICE = torch.device("cpu")
+
+# ── Ensemble configs (from optimizer) ──
 ENSEMBLE_CONFIGS = {
     "clean": {
         "models": ["image_ensemble_v2", "ai_gen", "spectral", "frequency"],
         "weights": [0.286, 0.357, 0.214, 0.143],
-        "threshold": 0.45,
+        "threshold": 0.45
     },
     "whatsapp": {
         "models": ["image_ensemble_v2"],
         "weights": [1.0],
-        "threshold": 0.65,
+        "threshold": 0.65
     },
     "instagram": {
         "models": ["image_ensemble_v2", "ai_gen"],
         "weights": [0.6, 0.4],
-        "threshold": 0.60,
+        "threshold": 0.60
     },
     "telegram": {
         "models": ["image_ensemble_v2", "ai_gen"],
         "weights": [0.75, 0.25],
-        "threshold": 0.60,
+        "threshold": 0.60
     },
     "screenshot": {
         "models": ["image_ensemble_v2", "spectral"],
         "weights": [0.667, 0.333],
-        "threshold": 0.45,
-    },
+        "threshold": 0.45
+    }
 }
 
-# ──────────────────────────────────────────────
-# Image models (lazy loaded)
-# ──────────────────────────────────────────────
-_models = {}
-IMG_SIZE = 256
-RADIAL_BINS = 128
-SPECTRUM_SIZE = 64
-
-def _get_transform():
-    """Standard transform for EfficientNet models."""
-    from torchvision import transforms
-    return transforms.Compose([
-        transforms.Resize((224, 224)),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-    ])
+# ── Standard transforms ──
+img_transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
+img_transform_300 = transforms.Compose([
+    transforms.Resize((300, 300)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
 
-_transform = None
-def get_transform():
-    global _transform
-    if _transform is None:
-        _transform = _get_transform()
-    return _transform
+# ── Model registry ──
+models = {}
 
 
+# ══════════════════════════════════════════════
+# IMAGE ENSEMBLE V2 (EfficientNet-B0, 16MB)
+# ══════════════════════════════════════════════
 def load_image_ensemble_v2():
-    if "image_ensemble_v2" in _models:
-        return _models["image_ensemble_v2"]
+    if "image_ensemble_v2" in models:
+        return
     print("Loading image_ensemble_v2 (EfficientNet-B0)...", flush=True)
     model = efficientnet_b0(weights=None)
     model.classifier[1] = nn.Linear(1280, 2)
-    ckpt = torch.load("image_ensemble_v2.pt", map_location="cpu", weights_only=False)
-    model.load_state_dict(ckpt)
+    sd = torch.load("image_ensemble_v2.pt", map_location=DEVICE, weights_only=False)
+    if isinstance(sd, dict) and "model_state_dict" in sd:
+        sd = sd["model_state_dict"]
+    model.load_state_dict(sd)
     model.eval()
-    _models["image_ensemble_v2"] = model
-    print("image_ensemble_v2 loaded.", flush=True)
-    return model
+    models["image_ensemble_v2"] = model
+    print("  ✓ image_ensemble_v2 loaded", flush=True)
 
 
+def infer_image_ensemble_v2(img: Image.Image) -> float:
+    load_image_ensemble_v2()
+    tensor = img_transform(img).unsqueeze(0).to(DEVICE)
+    with torch.no_grad():
+        logits = models["image_ensemble_v2"](tensor)
+        prob = torch.softmax(logits, dim=1)[0, 1].item()
+    return prob
+
+
+# ══════════════════════════════════════════════
+# AI GEN DETECTOR (EfficientNet-B3, 46MB)
+# ══════════════════════════════════════════════
 def load_ai_gen():
-    if "ai_gen" in _models:
-        return _models["ai_gen"]
-    print("Loading ai_gen_detector (EfficientNet-B3)...", flush=True)
+    if "ai_gen" in models:
+        return
+    print("Loading ai_gen (EfficientNet-B3)...", flush=True)
     model = efficientnet_b3(weights=None)
-    model.classifier[1] = nn.Linear(model.classifier[1].in_features, 2)
-    ckpt = torch.load("ai_gen_detector.pt", map_location="cpu", weights_only=False)
-    model.load_state_dict(ckpt["model_state_dict"])
+    model.classifier[1] = nn.Linear(1536, 2)
+    ckpt = torch.load("ai_gen_detector.pt", map_location=DEVICE, weights_only=False)
+    sd = ckpt.get("model_state_dict", ckpt)
+    model.load_state_dict(sd)
     model.eval()
-    _models["ai_gen"] = model
-    print("ai_gen loaded.", flush=True)
-    return model
-
-
-# ──────── Spectral model ────────
+    models["ai_gen"] = model
+    print("  ✓ ai_gen loaded", flush=True)
 
-class RadialCNN1D(nn.Module):
-    def __init__(self, input_len=RADIAL_BINS, embed_dim=128):
-        super().__init__()
-        self.net = nn.Sequential(
-            nn.Conv1d(1, 32, kernel_size=7, padding=3),
-            nn.BatchNorm1d(32), nn.ReLU(inplace=True), nn.MaxPool1d(2),
-            nn.Conv1d(32, 64, kernel_size=5, padding=2),
-            nn.BatchNorm1d(64), nn.ReLU(inplace=True), nn.MaxPool1d(2),
-            nn.Conv1d(64, 128, kernel_size=3, padding=1),
-            nn.BatchNorm1d(128), nn.ReLU(inplace=True), nn.AdaptiveAvgPool1d(4),
-        )
-        self.fc = nn.Linear(128 * 4, embed_dim)
 
-    def forward(self, x):
-        x = x.unsqueeze(1)
-        x = self.net(x)
-        x = x.flatten(1)
-        return self.fc(x)
+def infer_ai_gen(img: Image.Image) -> float:
+    load_ai_gen()
+    tensor = img_transform_300(img).unsqueeze(0).to(DEVICE)
+    with torch.no_grad():
+        logits = models["ai_gen"](tensor)
+        prob = torch.softmax(logits, dim=1)[0, 1].item()
+    return prob
 
 
-class DualStreamSpectralDetector(nn.Module):
-    def __init__(self, embed_dim=256, radial_bins=RADIAL_BINS):
+# ══════════════════════════════════════════════
+# SPECTRAL DETECTOR (Dual-stream ResNet, 137MB)
+# ══════════════════════════════════════════════
+class DualStreamSpectral(nn.Module):
+    """Dual-stream spectral detector: stream_a and stream_b on frequency features."""
+    def __init__(self):
         super().__init__()
-        r = resnet18(weights=None)
-        r.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
-        r.fc = nn.Linear(r.fc.in_features, embed_dim)
-        self.stream_a = r
-        self.stream_b = RadialCNN1D(input_len=radial_bins, embed_dim=embed_dim // 2)
-        fused_dim = embed_dim + embed_dim // 2
-        self.classifier = nn.Sequential(
-            nn.LayerNorm(fused_dim), nn.Dropout(0.3),
-            nn.Linear(fused_dim, 128), nn.ReLU(inplace=True),
-            nn.Dropout(0.2), nn.Linear(128, 1),
+        from torchvision.models import resnet34
+        self.stream_a = resnet34(weights=None)
+        self.stream_a.conv1 = nn.Conv2d(1, 64, 7, stride=2, padding=3, bias=False)
+        self.stream_b = resnet34(weights=None)
+        self.stream_b.conv1 = nn.Conv2d(1, 64, 7, stride=2, padding=3, bias=False)
+        feat_dim = 512 + 512
+        self.stream_a.fc = nn.Identity()
+        self.stream_b.fc = nn.Identity()
+        self.head = nn.Sequential(
+            nn.Linear(feat_dim, 256),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(256, 2)
         )
 
-    def forward(self, spectrum_img, radial):
-        a = self.stream_a(spectrum_img)
-        b = self.stream_b(radial)
-        fused = torch.cat([a, b], dim=1)
-        return self.classifier(fused).squeeze(-1)
-
-
-def compute_dct_spectrum(img_gray):
-    from scipy.fft import dctn
-    dct = dctn(img_gray.astype(np.float64), type=2, norm="ortho")
-    log_ps = np.log1p(dct ** 2)
-    mx = log_ps.max()
-    if mx > 0:
-        log_ps /= mx
-    return log_ps.astype(np.float32)
-
-
-def compute_radial_profile(img_gray, num_bins=RADIAL_BINS):
-    f = np.fft.fft2(img_gray.astype(np.float64))
-    f_shift = np.fft.fftshift(f)
-    magnitude = np.log1p(np.abs(f_shift))
-    h, w = magnitude.shape
-    cy, cx = h // 2, w // 2
-    y, x = np.ogrid[:h, :w]
-    r = np.sqrt((x - cx) ** 2 + (y - cy) ** 2)
-    max_r = np.sqrt(cx ** 2 + cy ** 2)
-    bin_edges = np.linspace(0, max_r, num_bins + 1)
-    profile = np.zeros(num_bins, dtype=np.float32)
-    for i in range(num_bins):
-        mask = (r >= bin_edges[i]) & (r < bin_edges[i + 1])
-        vals = magnitude[mask]
-        profile[i] = vals.mean() if vals.size > 0 else 0.0
-    mx = profile.max()
-    if mx > 0:
-        profile /= mx
-    return profile
+    def forward(self, dct_input, fft_input):
+        a = self.stream_a(dct_input)
+        b = self.stream_b(fft_input)
+        combined = torch.cat([a, b], dim=1)
+        return self.head(combined)
 
 
 def load_spectral():
-    if "spectral" in _models:
-        return _models["spectral"]
-    print("Loading spectral_detector (DualStream)...", flush=True)
-    model = DualStreamSpectralDetector()
-    ckpt = torch.load("spectral_detector.pt", map_location="cpu", weights_only=False)
-    model.load_state_dict(ckpt["model"])
+    if "spectral" in models:
+        return
+    print("Loading spectral detector...", flush=True)
+    model = DualStreamSpectral()
+    ckpt = torch.load("spectral_detector.pt", map_location=DEVICE, weights_only=False)
+    sd = ckpt.get("model", ckpt.get("model_state_dict", ckpt))
+    model.load_state_dict(sd, strict=False)
     model.eval()
-    _models["spectral"] = model
-    print("spectral loaded.", flush=True)
-    return model
-
-
-# ──────── Frequency analyzer model ────────
-
-BENFORD_EXPECTED = np.array([np.log10(1 + 1/d) for d in range(1, 10)])
-BENFORD_DIM = 18
-DCT_STATS_DIM = 10
-WAVELET_DIM = 26
-HANDCRAFTED_DIM = 54
-
-
-def extract_dct_benford_features(img_gray):
-    from scipy.fft import dctn
-    h, w = img_gray.shape
-    all_coeffs = []
-    block_size = 8
-    for i in range(0, h - block_size + 1, block_size):
-        for j in range(0, w - block_size + 1, block_size):
-            block = img_gray[i:i+block_size, j:j+block_size].astype(np.float64)
-            dct_block = dctn(block, type=2, norm='ortho')
-            coeffs = dct_block.flatten()[1:]
-            all_coeffs.extend(coeffs)
-    all_coeffs = np.abs(np.array(all_coeffs))
-    all_coeffs = all_coeffs[all_coeffs > 0]
-    if len(all_coeffs) < 100:
-        return np.zeros(18, dtype=np.float32)
-    first_digits = np.floor(all_coeffs / (10 ** np.floor(np.log10(all_coeffs)))).astype(int)
-    first_digits = first_digits[(first_digits >= 1) & (first_digits <= 9)]
-    if len(first_digits) < 100:
-        return np.zeros(18, dtype=np.float32)
-    observed = np.zeros(9, dtype=np.float64)
-    for d in range(1, 10):
-        observed[d-1] = np.sum(first_digits == d) / len(first_digits)
-    deviation = (observed - BENFORD_EXPECTED) ** 2 / (BENFORD_EXPECTED + 1e-10)
-    return np.concatenate([observed.astype(np.float32), deviation.astype(np.float32)])
-
-
-def extract_dct_block_stats(img_gray):
-    from scipy.fft import dctn
-    h, w = img_gray.shape
-    block_size = 8
-    energy_map = np.zeros((block_size, block_size), dtype=np.float64)
-    count = 0
-    for i in range(0, h - block_size + 1, block_size):
-        for j in range(0, w - block_size + 1, block_size):
-            block = img_gray[i:i+block_size, j:j+block_size].astype(np.float64)
-            dct_block = dctn(block, type=2, norm='ortho')
-            energy_map += dct_block ** 2
-            count += 1
-    if count > 0:
-        energy_map /= count
-    total = energy_map.sum() + 1e-10
-    energy_map /= total
-    low = energy_map[:4, :4].sum()
-    mid = energy_map.sum() - low - energy_map[4:, 4:].sum()
-    high = energy_map[4:, 4:].sum()
-    flat = energy_map.flatten()
-    return np.array([low, mid, high, flat.std(), np.median(flat),
-                     flat[0], flat[-1], np.log1p(flat.max()), np.log1p(flat.min()),
-                     np.sum(flat > 0.01)], dtype=np.float32)
-
-
-def extract_wavelet_features(img_gray):
-    import pywt
-    features = []
-    for wavelet in ['haar', 'db4']:
-        coeffs = pywt.wavedec2(img_gray.astype(np.float64), wavelet, level=3)
-        approx_energy = np.sum(coeffs[0] ** 2)
-        total_energy = approx_energy
-        detail_energies = []
-        for level_coeffs in coeffs[1:]:
-            level_energy = sum(np.sum(c ** 2) for c in level_coeffs)
-            detail_energies.append(level_energy)
-            total_energy += level_energy
-        total_energy = max(total_energy, 1e-10)
-        features.append(approx_energy / total_energy)
-        for de in detail_energies:
-            features.append(de / total_energy)
-        finest = coeffs[-1]
-        for c in finest:
-            features.extend([
-                np.std(c),
-                np.mean(np.abs(c)),
-                float(np.sum(np.abs(c) > np.std(c) * 3)) / max(c.size, 1),
-            ])
-    return np.array(features, dtype=np.float32)
-
-
-class FrequencyAnalyzerModel(nn.Module):
-    def __init__(self, handcrafted_dim=HANDCRAFTED_DIM, spectrum_size=SPECTRUM_SIZE):
-        super().__init__()
-        self.mlp = nn.Sequential(
-            nn.Linear(handcrafted_dim, 128), nn.BatchNorm1d(128),
-            nn.ReLU(inplace=True), nn.Dropout(0.3),
-            nn.Linear(128, 64), nn.ReLU(inplace=True),
-        )
-        self.cnn = nn.Sequential(
-            nn.Conv2d(1, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(inplace=True), nn.MaxPool2d(2),
-            nn.Conv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True), nn.MaxPool2d(2),
-            nn.Conv2d(64, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(inplace=True), nn.AdaptiveAvgPool2d(4),
-        )
-        self.cnn_fc = nn.Linear(128 * 4 * 4, 128)
-        self.classifier = nn.Sequential(
-            nn.LayerNorm(64 + 128), nn.Dropout(0.3),
-            nn.Linear(64 + 128, 64), nn.ReLU(inplace=True),
-            nn.Dropout(0.2), nn.Linear(64, 1),
-        )
+    models["spectral"] = model
+    print("  ✓ spectral loaded", flush=True)
 
-    def forward(self, handcrafted, spectrum):
-        h = self.mlp(handcrafted)
-        s = self.cnn(spectrum)
-        s = s.flatten(1)
-        s = self.cnn_fc(s)
-        fused = torch.cat([h, s], dim=1)
-        return self.classifier(fused).squeeze(-1)
+
+def extract_spectral_features(img: Image.Image):
+    """Convert image to DCT and FFT magnitude spectral maps."""
+    gray = np.array(img.convert("L").resize((224, 224)), dtype=np.float32) / 255.0
+    
+    # DCT via scipy if available, else approximate with FFT
+    try:
+        from scipy.fftpack import dct
+        dct_map = dct(dct(gray, axis=0, norm='ortho'), axis=1, norm='ortho')
+    except ImportError:
+        dct_map = np.real(np.fft.fft2(gray))
+    
+    # FFT magnitude spectrum
+    fft = np.fft.fft2(gray)
+    fft_shift = np.fft.fftshift(fft)
+    magnitude = np.log1p(np.abs(fft_shift))
+    
+    # Normalize
+    dct_norm = (dct_map - dct_map.mean()) / (dct_map.std() + 1e-8)
+    mag_norm = (magnitude - magnitude.mean()) / (magnitude.std() + 1e-8)
+    
+    dct_tensor = torch.from_numpy(dct_norm).unsqueeze(0).unsqueeze(0).float()
+    fft_tensor = torch.from_numpy(mag_norm).unsqueeze(0).unsqueeze(0).float()
+    return dct_tensor, fft_tensor
+
+
+def infer_spectral(img: Image.Image) -> float:
+    load_spectral()
+    dct_t, fft_t = extract_spectral_features(img)
+    with torch.no_grad():
+        logits = models["spectral"](dct_t.to(DEVICE), fft_t.to(DEVICE))
+        prob = torch.softmax(logits, dim=1)[0, 1].item()
+    return prob
 
 
-def extract_dct_spectrum_image(img_gray, size=SPECTRUM_SIZE):
-    from scipy.fft import dctn
-    dct = dctn(img_gray.astype(np.float64), type=2, norm='ortho')
-    log_ps = np.log1p(np.abs(dct))
-    mx = log_ps.max()
-    if mx > 0:
-        log_ps /= mx
-    img = Image.fromarray((log_ps * 255).astype(np.uint8))
-    img = img.resize((size, size), Image.LANCZOS)
-    return np.array(img, dtype=np.float32) / 255.0
+# ══════════════════════════════════════════════
+# FREQUENCY ANALYZER (MLP on handcrafted, 1.6MB)
+# ══════════════════════════════════════════════
+class FrequencyMLP(nn.Module):
+    def __init__(self, input_dim=54):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(input_dim, 128),
+            nn.BatchNorm1d(128),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(128, 64),
+            nn.BatchNorm1d(64),
+            nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(64, 2)
+        )
+    def forward(self, x):
+        return self.mlp(x)
 
 
 def load_frequency():
-    if "frequency" in _models:
-        return _models["frequency"]
-    print("Loading frequency_analyzer...", flush=True)
-    model = FrequencyAnalyzerModel()
-    ckpt = torch.load("frequency_analyzer.pt", map_location="cpu", weights_only=False)
-    model.load_state_dict(ckpt["model"])
+    if "frequency" in models:
+        return
+    print("Loading frequency analyzer...", flush=True)
+    ckpt = torch.load("frequency_analyzer.pt", map_location=DEVICE, weights_only=False)
+    dims = ckpt.get("feature_dims", {})
+    input_dim = dims.get("handcrafted_total", 54)
+    model = FrequencyMLP(input_dim)
+    sd = ckpt.get("model", ckpt.get("model_state_dict", ckpt))
+    model.load_state_dict(sd, strict=False)
     model.eval()
-    _models["frequency"] = model
-    print("frequency loaded.", flush=True)
-    return model
-
-
-# ──────── Model inference functions ────────
-
-def infer_efficientnet(model, img_pil):
-    """Run inference on an EfficientNet model, return fake probability."""
-    tensor = get_transform()(img_pil).unsqueeze(0)
-    with torch.no_grad():
-        logits = model(tensor)
-        probs = torch.softmax(logits, dim=-1)
-    return float(probs[0][1])  # fake class probability
-
-
-def infer_spectral(model, img_pil):
-    """Run spectral model on image."""
-    gray = np.array(img_pil.convert("L").resize((IMG_SIZE, IMG_SIZE), Image.LANCZOS))
-    spectrum = compute_dct_spectrum(gray)
-    radial = compute_radial_profile(gray)
-    spec_tensor = torch.from_numpy(spectrum).unsqueeze(0).unsqueeze(0)  # (1,1,H,W)
-    rad_tensor = torch.from_numpy(radial).unsqueeze(0)  # (1, bins)
-    with torch.no_grad():
-        logit = model(spec_tensor, rad_tensor)
-    return float(torch.sigmoid(logit).item())
+    models["frequency"] = model
+    models["frequency_dims"] = dims
+    print(f"  ✓ frequency loaded (input_dim={input_dim})", flush=True)
 
 
-def infer_frequency(model, img_pil):
-    """Run frequency analyzer on image."""
-    gray = np.array(img_pil.convert("L").resize((IMG_SIZE, IMG_SIZE), Image.LANCZOS))
-    benford = extract_dct_benford_features(gray)
-    dct_stats = extract_dct_block_stats(gray)
-    wavelet = extract_wavelet_features(gray)
-    handcrafted = np.concatenate([benford, dct_stats, wavelet])
-    spectrum_img = extract_dct_spectrum_image(gray)
-
-    hc_tensor = torch.from_numpy(handcrafted).unsqueeze(0)
-    sp_tensor = torch.from_numpy(spectrum_img).unsqueeze(0).unsqueeze(0)  # (1,1,64,64)
+def extract_frequency_features(img: Image.Image) -> torch.Tensor:
+    """Extract handcrafted frequency features: Benford, DCT stats, wavelet."""
+    gray = np.array(img.convert("L").resize((256, 256)), dtype=np.float32)
+    features = []
+    
+    # Benford's Law analysis on DCT coefficients
+    try:
+        from scipy.fftpack import dct
+        dct_coeffs = dct(dct(gray, axis=0, norm='ortho'), axis=1, norm='ortho').flatten()
+    except ImportError:
+        dct_coeffs = np.real(np.fft.fft2(gray)).flatten()
+    
+    abs_coeffs = np.abs(dct_coeffs[dct_coeffs != 0])
+    if len(abs_coeffs) > 0:
+        first_digits = (abs_coeffs / (10 ** np.floor(np.log10(abs_coeffs + 1e-10)))).astype(int)
+        first_digits = np.clip(first_digits, 1, 9)
+        benford_dist = np.zeros(9)
+        for d in range(1, 10):
+            benford_dist[d-1] = np.sum(first_digits == d) / len(first_digits)
+        benford_expected = np.log10(1 + 1/np.arange(1, 10))
+        features.extend(benford_dist.tolist())
+        features.append(float(np.sum(np.abs(benford_dist - benford_expected))))
+    else:
+        features.extend([0.0] * 10)
+    
+    # DCT statistics
+    dct_flat = dct_coeffs
+    dct_stats = [
+        float(np.mean(np.abs(dct_flat))),
+        float(np.std(dct_flat)),
+        float(np.median(np.abs(dct_flat))),
+        float(np.percentile(np.abs(dct_flat), 25)),
+        float(np.percentile(np.abs(dct_flat), 75)),
+        float(np.percentile(np.abs(dct_flat), 95)),
+        float(np.mean(dct_flat ** 2)),
+        float(np.max(np.abs(dct_flat))),
+    ]
+    # Energy in frequency bands
+    n = len(dct_flat)
+    band_size = n // 4
+    for i in range(4):
+        band = dct_flat[i*band_size:(i+1)*band_size]
+        dct_stats.append(float(np.mean(band ** 2)))
+    features.extend(dct_stats)
+    
+    # Wavelet-like features (Haar approximation)
+    h, w = gray.shape
+    ll = (gray[0::2, 0::2] + gray[1::2, 0::2] + gray[0::2, 1::2] + gray[1::2, 1::2]) / 4
+    lh = (gray[0::2, 0::2] - gray[1::2, 0::2] + gray[0::2, 1::2] - gray[1::2, 1::2]) / 4
+    hl = (gray[0::2, 0::2] + gray[1::2, 0::2] - gray[0::2, 1::2] - gray[1::2, 1::2]) / 4
+    hh = (gray[0::2, 0::2] - gray[1::2, 0::2] - gray[0::2, 1::2] + gray[1::2, 1::2]) / 4
+    for sub in [ll, lh, hl, hh]:
+        features.extend([
+            float(np.mean(np.abs(sub))),
+            float(np.std(sub)),
+            float(np.max(np.abs(sub))),
+            float(np.mean(sub ** 2)),
+        ])
+    
+    # FFT magnitude stats  
+    fft = np.fft.fft2(gray)
+    mag = np.abs(np.fft.fftshift(fft))
+    log_mag = np.log1p(mag)
+    features.extend([
+        float(np.mean(log_mag)),
+        float(np.std(log_mag)),
+        float(np.median(log_mag)),
+        float(np.max(log_mag)),
+    ])
+    
+    # Pad or truncate to expected dimension
+    dims = models.get("frequency_dims", {})
+    expected = dims.get("handcrafted_total", 54)
+    if len(features) < expected:
+        features.extend([0.0] * (expected - len(features)))
+    features = features[:expected]
+    
+    return torch.tensor([features], dtype=torch.float32)
+
+
+def infer_frequency(img: Image.Image) -> float:
+    load_frequency()
+    feat = extract_frequency_features(img).to(DEVICE)
     with torch.no_grad():
-        logit = model(hc_tensor, sp_tensor)
-    return float(torch.sigmoid(logit).item())
+        logits = models["frequency"](feat)
+        prob = torch.softmax(logits, dim=1)[0, 1].item()
+    return prob
 
 
-# Model loader dispatch
-MODEL_LOADERS = {
-    "image_ensemble_v2": load_image_ensemble_v2,
-    "ai_gen": load_ai_gen,
-    "spectral": load_spectral,
-    "frequency": load_frequency,
-}
-
-MODEL_INFER = {
-    "image_ensemble_v2": lambda img: infer_efficientnet(load_image_ensemble_v2(), img),
-    "ai_gen": lambda img: infer_efficientnet(load_ai_gen(), img),
-    "spectral": lambda img: infer_spectral(load_spectral(), img),
-    "frequency": lambda img: infer_frequency(load_frequency(), img),
-}
-
-
-# ──────────────────────────────────────────────
-# Audio Model V10
-# ──────────────────────────────────────────────
-audio_model_v10 = None
+# ══════════════════════════════════════════════
+# AUDIO MODEL (Wav2Vec2 + classifier)
+# ══════════════════════════════════════════════
+audio_model = None
 audio_feature_extractor = None
+audio_classifier = None
 
-class AudioModelV10(nn.Module):
-    """Wav2Vec2 backbone + classifier head for 3-class audio deepfake detection."""
-    def __init__(self, num_classes=3):
+class AudioClassifier(nn.Module):
+    def __init__(self, input_dim=768, hidden_dim=256):
         super().__init__()
-        from transformers import Wav2Vec2Model
-        self.backbone = Wav2Vec2Model.from_pretrained("facebook/wav2vec2-base")
-        self.head = nn.Sequential(
-            nn.LayerNorm(768),
-            nn.Linear(768, 256),
-            nn.ReLU(inplace=True),
+        self.classifier = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            nn.ReLU(),
             nn.Dropout(0.3),
-            nn.Linear(256, 128),
-            nn.ReLU(inplace=True),
+            nn.Linear(hidden_dim, 64),
+            nn.ReLU(),
             nn.Dropout(0.2),
-            nn.Linear(128, num_classes),
+            nn.Linear(64, 2)
         )
-
-    def forward(self, input_values):
-        outputs = self.backbone(input_values)
-        hidden = outputs.last_hidden_state.mean(dim=1)
-        return self.head(hidden)
+    def forward(self, x):
+        return self.classifier(x)
 
 
-def load_audio_v10():
-    global audio_model_v10, audio_feature_extractor
-    if audio_model_v10 is not None:
+def load_audio():
+    global audio_model, audio_feature_extractor, audio_classifier
+    if audio_classifier is not None:
         return
-    print("Loading Audio V10 model...", flush=True)
-    from transformers import Wav2Vec2FeatureExtractor
+    print("Loading Wav2Vec2 + audio classifier...", flush=True)
     audio_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("facebook/wav2vec2-base")
-    audio_model_v10 = AudioModelV10(num_classes=3)
-    ckpt = torch.load("audio_deepfake_v10.pt", map_location="cpu", weights_only=False)
-    audio_model_v10.load_state_dict(ckpt["model_state_dict"])
-    audio_model_v10.eval()
-    print("Audio V10 loaded.", flush=True)
-
-
-# ──────────────────────────────────────────────
-# CLIP (legacy, for /video endpoint)
-# ──────────────────────────────────────────────
+    audio_model = Wav2Vec2Model.from_pretrained("facebook/wav2vec2-base")
+    audio_model.eval()
+    audio_classifier = AudioClassifier()
+    state = torch.load("audio_deepfake_model.pt", map_location="cpu", weights_only=False)
+    if isinstance(state, dict) and "classifier_state_dict" in state:
+        sd = state["classifier_state_dict"]
+    else:
+        sd = state
+    if isinstance(sd, dict) and any(k[0].isdigit() for k in sd.keys()):
+        sd = {f"classifier.{k}": v for k, v in sd.items()}
+    audio_classifier.load_state_dict(sd)
+    audio_classifier.eval()
+    print("  ✓ audio loaded", flush=True)
+
+
+# ══════════════════════════════════════════════
+# LEGACY CLIP PROBE (for /video frame analysis)
+# ══════════════════════════════════════════════
 clip_model = None
 clip_preprocess = None
 clip_probe = None
@@ -442,82 +381,95 @@ def load_clip():
     if clip_model is not None:
         return
     if not os.path.exists("clip_deepfake_probe.pt"):
-        print("CLIP probe not found, skipping.", flush=True)
+        print("CLIP probe not found, skipping", flush=True)
         return
     print("Loading CLIP ViT-L/14...", flush=True)
     import open_clip
     clip_model, _, clip_preprocess = open_clip.create_model_and_transforms('ViT-L-14', pretrained='openai')
     clip_model.eval()
     clip_probe = torch.load("clip_deepfake_probe.pt", map_location="cpu", weights_only=False)
-    print(f"CLIP probe loaded. Val accuracy: {clip_probe['val_accuracy']:.4f}", flush=True)
+    print(f"  ✓ CLIP probe loaded (val_acc={clip_probe.get('val_accuracy', 'N/A')})", flush=True)
 
 
-# ──────────────────────────────────────────────
-# Endpoints
-# ──────────────���───────────────────────────────
+# ══════════════════════════════════════════════
+# ENDPOINTS
+# ══════════════════════════════════════════════
 
 @app.get("/health")
 def health():
     return {
         "status": "ok",
-        "version": "2.0.0-v10",
-        "models": ["image_ensemble_v2", "ai_gen", "spectral", "frequency", "audio_v10", "clip-legacy"],
-        "endpoints": ["/image", "/audio", "/video"],
+        "version": "10.0.0",
+        "models": list(models.keys()),
         "platforms": list(ENSEMBLE_CONFIGS.keys()),
+        "endpoints": ["/image", "/audio", "/video"]
     }
 
 
 @app.post("/image")
 async def analyze_image(
     file: UploadFile = File(...),
-    platform: str = Query("clean", description="Platform: clean, whatsapp, instagram, telegram, screenshot"),
+    platform: str = Query("clean", regex="^(clean|whatsapp|instagram|telegram|screenshot)$")
 ):
-    """Analyze image with V10 ensemble, platform-aware weighting."""
+    """Analyze image using V10 ensemble with platform-specific config."""
     try:
-        platform = platform.lower()
-        if platform not in ENSEMBLE_CONFIGS:
-            platform = "clean"
-        config = ENSEMBLE_CONFIGS[platform]
-
         data = await file.read()
         img = Image.open(io.BytesIO(data)).convert("RGB")
-
+        
+        config = ENSEMBLE_CONFIGS[platform]
+        model_names = config["models"]
+        weights = config["weights"]
+        threshold = config["threshold"]
+        
         # Run each model in the ensemble
         ensemble_scores = {}
-        for model_name in config["models"]:
-            if model_name not in MODEL_INFER:
-                continue
-            score = MODEL_INFER[model_name](img)
-            ensemble_scores[model_name] = round(score, 4)
-
-        # Weighted combination
-        weighted_score = 0.0
-        total_weight = 0.0
-        for model_name, weight in zip(config["models"], config["weights"]):
-            if model_name in ensemble_scores:
-                weighted_score += weight * ensemble_scores[model_name]
-                total_weight += weight
-
+        infer_fns = {
+            "image_ensemble_v2": infer_image_ensemble_v2,
+            "ai_gen": infer_ai_gen,
+            "spectral": infer_spectral,
+            "frequency": infer_frequency,
+        }
+        
+        for name in model_names:
+            fn = infer_fns.get(name)
+            if fn:
+                try:
+                    ensemble_scores[name] = fn(img)
+                except Exception as e:
+                    print(f"Model {name} failed: {e}", flush=True)
+                    traceback.print_exc()
+                    ensemble_scores[name] = 0.5  # neutral fallback
+        
+        # Weighted average
+        total_weight = 0
+        weighted_score = 0
+        for name, w in zip(model_names, weights):
+            if name in ensemble_scores:
+                weighted_score += ensemble_scores[name] * w
+                total_weight += w
+        
         if total_weight > 0:
-            weighted_score /= total_weight
-
-        threshold = config["threshold"]
-        verdict = "fake" if weighted_score > threshold else "real"
-        fake_prob = weighted_score
-        real_prob = 1.0 - fake_prob
-
+            fake_prob = weighted_score / total_weight
+        else:
+            fake_prob = 0.5
+        
+        real_prob = 1 - fake_prob
+        verdict = "fake" if fake_prob > threshold else "real"
+        confidence = max(fake_prob, real_prob)
+        
         return {
             "verdict": verdict,
-            "confidence": round(max(fake_prob, real_prob), 4),
+            "confidence": round(confidence, 4),
             "scores": {
                 "real": round(real_prob, 4),
-                "fake": round(fake_prob, 4),
+                "fake": round(fake_prob, 4)
             },
-            "ensemble_scores": ensemble_scores,
+            "ensemble_scores": {k: round(v, 4) for k, v in ensemble_scores.items()},
             "platform": platform,
-            "models_used": config["models"],
+            "models_used": model_names,
             "threshold": threshold,
             "model": "kaeva-v10-ensemble",
+            "version": "10.0.0"
         }
     except Exception as e:
         traceback.print_exc()
@@ -526,42 +478,35 @@ async def analyze_image(
 
 @app.post("/audio")
 async def analyze_audio(file: UploadFile = File(...)):
-    """Analyze audio for deepfake (3-class: real/tts/vc)."""
+    """Analyze audio for deepfake indicators using Wav2Vec2 probe."""
     try:
-        load_audio_v10()
-        import librosa
+        load_audio()
         data = await file.read()
         audio_np, sr = librosa.load(io.BytesIO(data), sr=16000, mono=True)
-
-        # Truncate/pad to 10 seconds
-        max_len = 16000 * 10
+        max_len = 16000 * 5
         if len(audio_np) > max_len:
             audio_np = audio_np[:max_len]
-        elif len(audio_np) < 16000:
-            audio_np = np.pad(audio_np, (0, 16000 - len(audio_np)))
-
+        elif len(audio_np) < max_len:
+            audio_np = np.pad(audio_np, (0, max_len - len(audio_np)))
+        
         inputs = audio_feature_extractor(audio_np, sampling_rate=16000, return_tensors="pt", padding=True)
-
         with torch.no_grad():
-            logits = audio_model_v10(inputs["input_values"])
+            outputs = audio_model(**inputs)
+            hidden = outputs.last_hidden_state.mean(dim=1)
+            logits = audio_classifier(hidden)
             probs = torch.softmax(logits, dim=-1)
-
-        classes = ["real", "tts", "vc"]
-        class_probs = {c: round(float(probs[0][i]), 4) for i, c in enumerate(classes)}
-        pred_idx = int(probs[0].argmax())
-        verdict = "real" if pred_idx == 0 else "fake"
-        fake_type = classes[pred_idx] if pred_idx > 0 else None
-
+        
+        real_prob = float(probs[0][0])
+        fake_prob = float(probs[0][1])
+        verdict = "fake" if fake_prob > 0.5 else "real"
+        confidence = max(real_prob, fake_prob)
+        
         return {
             "verdict": verdict,
-            "confidence": round(float(probs[0][pred_idx]), 4),
-            "scores": {
-                "real": class_probs["real"],
-                "fake": round(1.0 - class_probs["real"], 4),
-            },
-            "class_scores": class_probs,
-            "fake_type": fake_type,
-            "model": "wav2vec2-v10-3class",
+            "confidence": round(confidence, 4),
+            "scores": {"real": round(real_prob, 4), "fake": round(fake_prob, 4)},
+            "model": "wav2vec2-audio-probe",
+            "version": "10.0.0"
         }
     except Exception as e:
         traceback.print_exc()
@@ -571,30 +516,26 @@ async def analyze_audio(file: UploadFile = File(...)):
 @app.post("/video")
 async def analyze_video(
     file: UploadFile = File(...),
-    platform: str = Query("clean", description="Platform for frame analysis"),
+    platform: str = Query("clean", regex="^(clean|whatsapp|instagram|telegram|screenshot)$")
 ):
-    """Analyze video: extract frames → V10 ensemble per frame, extract audio → audio model."""
+    """Analyze video: extract frames → V10 ensemble, extract audio → audio probe."""
     start_time = time.time()
     try:
         data = await file.read()
-
         with tempfile.TemporaryDirectory() as tmpdir:
             video_path = os.path.join(tmpdir, "input_video")
             with open(video_path, "wb") as f:
                 f.write(data)
-
+            
             # Get video info
-            probe_cmd = [
-                "ffprobe", "-v", "quiet", "-print_format", "json",
-                "-show_format", "-show_streams", video_path,
-            ]
+            probe_cmd = ["ffprobe", "-v", "quiet", "-print_format", "json", "-show_format", "-show_streams", video_path]
             probe_result = subprocess.run(probe_cmd, capture_output=True, text=True, timeout=15)
             video_info = json.loads(probe_result.stdout) if probe_result.returncode == 0 else {}
-
+            
             duration = float(video_info.get("format", {}).get("duration", 0))
             resolution = "unknown"
             fps = 30.0
-            has_audio_stream = False
+            has_audio = False
             for stream in video_info.get("streams", []):
                 if stream.get("codec_type") == "video":
                     resolution = f"{stream.get('width', '?')}x{stream.get('height', '?')}"
@@ -602,97 +543,95 @@ async def analyze_video(
                         num, den = stream.get("r_frame_rate", "30/1").split("/")
                         fps = float(num) / float(den)
                     except:
-                        fps = 30.0
+                        pass
                 elif stream.get("codec_type") == "audio":
-                    has_audio_stream = True
-
+                    has_audio = True
+            
             total_frames = int(duration * fps) if duration > 0 else 0
-
-            # Extract keyframes
             max_frames = min(8, max(1, int(duration))) if duration > 0 else 5
             frame_interval = max(1.0, duration / max_frames) if duration > 0 else 1.0
             frame_dir = os.path.join(tmpdir, "frames")
             os.makedirs(frame_dir)
-
+            
             ffmpeg_cmd = [
                 "ffmpeg", "-i", video_path,
                 "-vf", f"fps=1/{frame_interval}",
                 "-frames:v", str(max_frames),
                 "-q:v", "2",
                 os.path.join(frame_dir, "frame_%03d.jpg"),
-                "-y", "-loglevel", "error",
+                "-y", "-loglevel", "error"
             ]
             subprocess.run(ffmpeg_cmd, timeout=30, check=True)
-
+            
             # Run V10 ensemble on each frame
-            platform_lower = platform.lower()
-            if platform_lower not in ENSEMBLE_CONFIGS:
-                platform_lower = "clean"
-            config = ENSEMBLE_CONFIGS[platform_lower]
-
             frame_files = sorted([f for f in os.listdir(frame_dir) if f.endswith(".jpg")])
             frame_scores = []
+            config = ENSEMBLE_CONFIGS[platform]
+            
             for fname in frame_files:
                 fpath = os.path.join(frame_dir, fname)
                 img = Image.open(fpath).convert("RGB")
+                
                 scores = {}
-                for model_name in config["models"]:
-                    if model_name in MODEL_INFER:
-                        scores[model_name] = MODEL_INFER[model_name](img)
-                weighted = sum(
-                    w * scores.get(m, 0) for m, w in zip(config["models"], config["weights"])
-                )
-                tw = sum(w for m, w in zip(config["models"], config["weights"]) if m in scores)
-                if tw > 0:
-                    weighted /= tw
-                frame_scores.append(round(weighted, 4))
-
+                infer_fns = {
+                    "image_ensemble_v2": infer_image_ensemble_v2,
+                    "ai_gen": infer_ai_gen,
+                    "spectral": infer_spectral,
+                    "frequency": infer_frequency,
+                }
+                for name in config["models"]:
+                    fn = infer_fns.get(name)
+                    if fn:
+                        try:
+                            scores[name] = fn(img)
+                        except:
+                            scores[name] = 0.5
+                
+                weighted = sum(scores.get(n, 0.5) * w for n, w in zip(config["models"], config["weights"]))
+                total_w = sum(w for n, w in zip(config["models"], config["weights"]) if n in scores)
+                frame_score = weighted / total_w if total_w > 0 else 0.5
+                frame_scores.append(round(frame_score, 4))
+            
             temporal_consistency = 1.0 - float(np.std(frame_scores)) if len(frame_scores) > 1 else 1.0
-
+            avg_frame_score = float(np.mean(frame_scores)) if frame_scores else 0.5
+            
             # Audio analysis
             audio_verdict = None
             audio_confidence = None
             audio_scores = None
-            audio_class_scores = None
-            if has_audio_stream:
+            if has_audio:
                 audio_path = os.path.join(tmpdir, "audio.wav")
-                audio_cmd = [
-                    "ffmpeg", "-i", video_path,
-                    "-vn", "-acodec", "pcm_s16le", "-ar", "16000", "-ac", "1",
-                    audio_path, "-y", "-loglevel", "error",
-                ]
+                audio_cmd = ["ffmpeg", "-i", video_path, "-vn", "-acodec", "pcm_s16le", "-ar", "16000", "-ac", "1", audio_path, "-y", "-loglevel", "error"]
                 audio_extract = subprocess.run(audio_cmd, timeout=20)
                 if audio_extract.returncode == 0 and os.path.exists(audio_path):
                     try:
-                        load_audio_v10()
-                        import librosa
+                        load_audio()
                         audio_np, sr = librosa.load(audio_path, sr=16000, mono=True)
                         max_audio_len = 16000 * 10
                         if len(audio_np) > max_audio_len:
                             audio_np = audio_np[:max_audio_len]
                         elif len(audio_np) < 16000:
                             audio_np = np.pad(audio_np, (0, 16000 - len(audio_np)))
-
                         inputs = audio_feature_extractor(audio_np, sampling_rate=16000, return_tensors="pt", padding=True)
                         with torch.no_grad():
-                            logits = audio_model_v10(inputs["input_values"])
+                            outputs = audio_model(**inputs)
+                            hidden = outputs.last_hidden_state.mean(dim=1)
+                            logits = audio_classifier(hidden)
                             probs = torch.softmax(logits, dim=-1)
-
-                        classes = ["real", "tts", "vc"]
-                        audio_class_scores = {c: round(float(probs[0][i]), 4) for i, c in enumerate(classes)}
-                        pred_idx = int(probs[0].argmax())
-                        audio_verdict = "real" if pred_idx == 0 else "fake"
-                        audio_confidence = round(float(probs[0][pred_idx]), 4)
-                        audio_scores = {"real": audio_class_scores["real"], "fake": round(1.0 - audio_class_scores["real"], 4)}
+                        rp = float(probs[0][0])
+                        fp = float(probs[0][1])
+                        audio_verdict = "fake" if fp > 0.5 else "real"
+                        audio_confidence = round(max(rp, fp), 4)
+                        audio_scores = {"real": round(rp, 4), "fake": round(fp, 4)}
                     except Exception as ae:
                         print(f"Audio analysis error: {ae}", flush=True)
-
-            avg_frame_score = float(np.mean(frame_scores)) if frame_scores else 0.5
+            
+            # Overall score
             if audio_scores:
                 overall_score = avg_frame_score * 0.7 + audio_scores["fake"] * 0.3
             else:
                 overall_score = avg_frame_score
-
+            
             flags = []
             if avg_frame_score > 0.7:
                 flags.append("HIGH_FAKE_SCORE_ACROSS_FRAMES")
@@ -702,10 +641,9 @@ async def analyze_video(
                 flags.append("AUDIO_FAKE_DETECTED")
             if audio_scores and ((avg_frame_score > 0.5) != (audio_scores["fake"] > 0.5)):
                 flags.append("AUDIO_VISUAL_DISAGREEMENT")
-
+            
             verdict = "fake" if overall_score > config["threshold"] else "real"
-            processing_time_ms = int((time.time() - start_time) * 1000)
-
+            
             return {
                 "frame_scores": frame_scores,
                 "temporal_consistency": round(temporal_consistency, 4),
@@ -716,13 +654,13 @@ async def analyze_video(
                 "fps": round(fps, 2),
                 "resolution": resolution,
                 "flags": flags,
-                "processing_time_ms": processing_time_ms,
+                "processing_time_ms": int((time.time() - start_time) * 1000),
                 "audio_verdict": audio_verdict,
                 "audio_confidence": audio_confidence,
                 "audio_scores": audio_scores,
-                "audio_class_scores": audio_class_scores,
-                "platform": platform_lower,
-                "model": "kaeva-v10-ensemble + wav2vec2-v10",
+                "platform": platform,
+                "model": "kaeva-v10-ensemble",
+                "version": "10.0.0"
             }
     except subprocess.TimeoutExpired:
         raise HTTPException(504, "Video processing timed out")