helo-ayush/Diarization_VoiceFingerprinted

1# ==============================================================================2# VOICE FINGERPRINT PROCESSOR3# Extracts 256-dimensional math vectors directly from raw audio waves.4# Uses Pyannote (WeSpeaker ResNet) to uniquely identify human voices.5# ==============================================================================6import io7import os8import tempfile9import torch10import torchaudio11import subprocess12 13print("⏳ Loading Pyannote Voice Fingerprinting model...")14 15_model = None16 17def _get_model():18    global _model19    if _model is None:20        from pyannote.audio import Model21        from pyannote.audio import Inference22        import torch23 24        import torchaudio25        import soundfile as sf26        import warnings27 28        # PyTorch 2.6+ defaults to weights_only=True, which breaks Pyannote.29        # We temporarily patch torch.load to use weights_only=False.30        _orig_load = torch.load31        def _safe_load(*args, **kwargs):32            kwargs['weights_only'] = False33            return _orig_load(*args, **kwargs)34        35        torch.load = _safe_load36 37        # Patch torchaudio.load to use soundfile instead of torchcodec (which fails on Windows)38        if not hasattr(torchaudio, "_patched"):39            def _sf_load(uri, **kwargs):40                data, samplerate = sf.read(uri, dtype='float32')41                if data.ndim == 1:42                    data = data.reshape(-1, 1)43                # soundfile returns (frames, channels), torchaudio expects (channels, frames)44                tensor = torch.from_numpy(data.T)45                return tensor, samplerate46                47            torchaudio.load = _sf_load48 49            # Patch torchaudio.info which was completely removed in torchaudio 2.1050            class MockTorchaudioInfo:51                def __init__(self, uri):52                    info = sf.info(uri)53                    self.num_frames = info.frames54                    self.sample_rate = info.samplerate55            56            torchaudio.info = MockTorchaudioInfo57            torchaudio._patched = True58 59        # Get HF Token from environment (.env)60        hf_token = os.environ.get("HF_TOKEN")61 62        # We use a public pyannote embedding model that doesn't strictly 63        # require an HF token for this specific architecture.64        # It yields a 256-dimensional embedding.65        try:66            model = Model.from_pretrained(67                "pyannote/wespeaker-voxceleb-resnet34-LM",68                use_auth_token=hf_token69            )70            _model = Inference(model, window="whole")71            print("✅ Pyannote Voice Fingerprinting loaded.")72        except Exception as e:73            print(f"❌ Failed to load Pyannote preset. You may need an HF_TOKEN: {e}")74            raise e75        finally:76            # Restore original torch.load77            torch.load = _orig_load78 79    return _model80 81# Try to load eagerly, but don't crash the server if it fails82try:83    _get_model()84except Exception as e:85    print(f"⚠️ Pyannote model will lazy-load on first use: {e}")86 87 88def extract_voice_embedding(audio_bytes: bytes, original_filename: str, start_sec: float = None, end_sec: float = None) -> list[float]:89    """90    Given a short voice clip, extracts a robust 256-dimensional Voice Print (Embedding).91    92    1. Converts whatever audio format to 16kHz Mono WAV using FFmpeg93    2. Optional: Crops the audio using start_sec and end_sec if provided94    3. Passes it through Pyannote's WeSpeaker ResNet95    4. Returns the mathematical embedding as a float array96    """97    inference = _get_model()98 99    ext = os.path.splitext(original_filename)[1] or ".input"100    with tempfile.NamedTemporaryFile(suffix=ext, delete=False) as tmp_in:101        tmp_in.write(audio_bytes)102        tmp_in_path = tmp_in.name103 104    tmp_wav_path = tmp_in_path + ".wav"105 106    try:107        # Step 1: Force standard format (16kHz, mono) so the AI reads it correctly108        ffmpeg_cmd = ["ffmpeg", "-y"]109        if start_sec is not None and end_sec is not None:110            ffmpeg_cmd.extend(["-ss", str(start_sec), "-to", str(end_sec)])111        ffmpeg_cmd.extend(["-i", tmp_in_path, "-ac", "1", "-ar", "16000", "-sample_fmt", "s16", tmp_wav_path])112        113        subprocess.run(114            ffmpeg_cmd,115            capture_output=True,116            check=True,117        )118 119        # Step 2: Extract embeddings120        # Pyannote Inference(window="whole") returns a numpy array representing the whole file121        embedding_ndarray = inference(tmp_wav_path)122        123        # Convert numpy array to flat float list124        flat_embedding = embedding_ndarray.flatten().tolist()125        126        return flat_embedding127 128    finally:129        for path in [tmp_in_path, tmp_wav_path]:130            try:131                os.unlink(path)132            except OSError:133                pass134