Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn.functional as F | |
| import numpy as np | |
| from transformers import AutoFeatureExtractor, AutoModelForAudioClassification | |
| from app.config import settings | |
| import logging | |
| import gc | |
| # Setup logging | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger(__name__) | |
| class VoiceDetector: | |
| _instance = None | |
| def __new__(cls): | |
| if cls._instance is None: | |
| cls._instance = super(VoiceDetector, cls).__new__(cls) | |
| cls._instance.model = None | |
| cls._instance.feature_extractor = None | |
| # Force CPU to save memory on free tier | |
| cls._instance.device = "cpu" | |
| cls._instance.load_model() | |
| return cls._instance | |
| def load_model(self): | |
| try: | |
| logger.info(f"Loading model {settings.MODEL_NAME} on {self.device}...") | |
| # Clear memory before loading | |
| gc.collect() | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| # Load with memory optimization | |
| self.feature_extractor = AutoFeatureExtractor.from_pretrained( | |
| settings.MODEL_NAME | |
| ) | |
| self.model = AutoModelForAudioClassification.from_pretrained( | |
| settings.MODEL_NAME, | |
| low_cpu_mem_usage=True, # Memory optimization | |
| torch_dtype=torch.float32 | |
| ) | |
| self.model.to(self.device) | |
| self.model.eval() | |
| # Clear unused memory | |
| gc.collect() | |
| logger.info("Model loaded successfully.") | |
| except Exception as e: | |
| logger.error(f"Failed to load model: {e}") | |
| raise RuntimeError(f"Failed to load model: {e}") | |
| def calibrate_confidence(self, probs, temperature=1.5): | |
| """ | |
| Apply temperature scaling to calibrate confidence scores. | |
| This makes the model less overconfident and more reliable. | |
| Temperature > 1.0 makes predictions less confident (more realistic) | |
| Temperature < 1.0 makes predictions more confident | |
| """ | |
| # Apply temperature scaling to logits before softmax | |
| logits = torch.log(probs + 1e-10) # Convert back to logits | |
| scaled_logits = logits / temperature | |
| calibrated_probs = F.softmax(scaled_logits, dim=-1) | |
| return calibrated_probs | |
| def predict(self, audio_array): | |
| """ | |
| Refined prediction for stability. | |
| """ | |
| if self.model is None: | |
| self.load_model() | |
| try: | |
| # Prepare input | |
| inputs = self.feature_extractor( | |
| audio_array, | |
| sampling_rate=settings.SAMPLE_RATE, | |
| return_tensors="pt", | |
| padding=True | |
| ) | |
| inputs = {key: val.to(self.device) for key, val in inputs.items()} | |
| # Inference | |
| with torch.no_grad(): | |
| logits = self.model(**inputs).logits | |
| # Use raw softmax for the base confidence | |
| probs = F.softmax(logits, dim=-1) | |
| # Get model labels from config | |
| id2label = self.model.config.id2label | |
| # Get the predicted class index | |
| pred_idx = torch.argmax(probs, dim=-1).item() | |
| label = str(id2label[pred_idx]).lower() | |
| confidence = probs[0][pred_idx].item() | |
| logger.info(f"Model Raw Output: Index={pred_idx}, Label={label}, Confidence={confidence:.4f}") | |
| # Robust Mapping Logic | |
| # mo-thecreator/Deepfake-audio-detection usually uses: | |
| # 0 -> REAL, 1 -> FAKE | |
| is_ai = False | |
| if "fake" in label or "spoof" in label: | |
| is_ai = True | |
| elif "real" in label or "bonafide" in label: | |
| is_ai = False | |
| else: | |
| # Direct index mapping fallback (very safe for this specific model) | |
| if pred_idx == 1: | |
| is_ai = True | |
| else: | |
| is_ai = False | |
| result_label = "AI_GENERATED" if is_ai else "HUMAN" | |
| # Stability check: If confidence is too low (< 0.6), | |
| # the model is essentially guessing. | |
| if confidence < 0.6: | |
| logger.info(f"Low confidence ({confidence:.4f}) detected. Result might be uncertain.") | |
| return result_label, confidence | |
| except Exception as e: | |
| logger.error(f"Prediction error: {e}") | |
| raise RuntimeError(f"Prediction failed: {e}") | |
| voice_detector = VoiceDetector() | |