import json import argparse import torch import numpy as np from pathlib import Path from tqdm import tqdm from pyannote.audio import Model, Inference from sklearn.cluster import AgglomerativeClustering # import umap.umap_ as umap # Optional for visualizing, skipping for speed """ tools/cluster_with_pyannote.py WHAT THIS SCRIPT DOES: ---------------------- 1. LOADS DATA: Reads your 'train_manifest.jsonl' to find all your audio files. 2. EXTRACTS "FINGERPRINTS": Uses the `pyannote/embedding` model to listen to each file and convert the voice into a mathematical vector (embedding). This is far more accurate than the previous Method. 3. CLUSTERS THEM: Instead of just comparing A vs B, it collects ALL fingerprints and groups them into "Families" (Speakers) using Agglomerative Clustering. 4. SAVES RESULT: Writes a new manifest with corrected 'speaker' IDs (e.g., 'speaker_0', 'speaker_1'). REQUIREMENTS: ------------- - pip install pyannote.audio sklearn - HuggingFace Token (exported as env var `HF_TOKEN` or logged in via `huggingface-cli login`) - Accepted license for `pyannote/embedding` on HuggingFace Hub. """ def parse_args(): parser = argparse.ArgumentParser(description="Cluster speakers using Pyannote embeddings.") parser.add_argument("--input", type=Path, required=True, help="Input manifest path.") parser.add_argument("--output", type=Path, required=True, help="Output manifest path.") parser.add_argument("--token", type=str, default=None, help="HuggingFace Token (optional if logged in).") parser.add_argument("--batch-size", type=int, default=32, help="Inference batch size.") parser.add_argument("--num-clusters", type=int, default=20000, help="Number of speakers to group into. Higher = stricter (more speakers). Lower = looser (fewer).") parser.add_argument("--detect-multi-speaker", action="store_true", help="Detect files with >1 speaker using Diarization pipeline (WARNING: Very Slow!).") return parser.parse_args() def main(): args = parse_args() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") # 1. Load Model print("Loading Pyannote Embedding Model (Wespeaker)...") try: # Fix for PyTorch 2.6+ safe globals issue try: import pyannote.audio.core.task torch.serialization.add_safe_globals([ torch.torch_version.TorchVersion, pyannote.audio.core.task.Specifications, pyannote.audio.core.task.Problem, pyannote.audio.core.task.Resolution, ]) except (AttributeError, ImportError): pass # Older torch versions or missing imports # Using the modern Wespeaker model which is the successor to pyannote/embedding model = Model.from_pretrained("pyannote/wespeaker-voxceleb-resnet34-LM", use_auth_token=args.token) inference = Inference(model, window="whole", device=device) except Exception as e: print(f"\n[ERROR] Could not load Pyannote model: {e}") print("Make sure you accepted the license at: https://huggingface.co/pyannote/wespeaker-voxceleb-resnet34-LM") print("And ran: huggingface-cli login\n") return # Load Diarization Pipeline ONLY if requested (very slow!) diarization_pipeline = None if args.detect_multi_speaker: print("\n[WARNING] Multi-speaker detection is ENABLED. This will drastically increase processing time!") print("Loading Pyannote Diarization Pipeline (Offline Cache)...") from pyannote.audio import Pipeline import os # Point directly to the cached config file to bypass auth errors cache_dir = Path(os.path.expanduser("~")) / ".cache/huggingface/hub/models--pyannote--speaker-diarization-3.1/snapshots" try: # Find the first snapshot folder if cache_dir.exists(): snap_dirs = [d for d in cache_dir.iterdir() if d.is_dir()] if snap_dirs: config_path = snap_dirs[0] / "config.yaml" diarization_pipeline = Pipeline.from_pretrained(str(config_path)) diarization_pipeline.to(device) print("Diarization Pipeline loaded successfully from cache.") else: print("[ERROR] No snapshot found in diarization cache.") else: print("[ERROR] Diarization cache directory not found.") except Exception as e: print(f"[ERROR] Could not load Diarization Pipeline: {e}") return # 2. Read Manifest print(f"Reading manifest: {args.input}") entries = [] with args.input.open("r", encoding="utf-8") as f: for line in f: if line.strip(): entries.append(json.loads(line)) # 3. Extract Embeddings (with caching) cache_path = args.input.parent / f"{args.input.stem}_embeddings.npy" cache_ids_path = args.input.parent / f"{args.input.stem}_ids.json" if cache_path.exists() and cache_ids_path.exists(): print(f"Loading cached embeddings from {cache_path}...") embeddings = np.load(cache_path) with open(cache_ids_path, 'r', encoding='utf-8') as f: valid_indices = json.load(f) # Filter entries to match the cache # We need to ensure 'entries' aligns with 'embeddings' # The cache stores indices of the ORIGINAL entries list that were valid valid_entries = [entries[i] for i in valid_indices] else: print(f"Extracting embeddings for {len(entries)} files...") embeddings = [] valid_entries = [] valid_indices = [] multi_speaker_files = [] # Resolving relative paths dataset_root = args.input.parent.parent for i, entry in enumerate(tqdm(entries)): audio_path = Path(entry['audio_path']) # FIX: The audio is likely in the 'dataset' subfolder possible_paths = [ audio_path, Path("dataset") / audio_path, args.input.parent / audio_path, Path("thaitts_emilia_dataset") / audio_path.name ] found_path = None for p in possible_paths: if p.exists(): found_path = p break if not found_path: continue try: embedding = inference(str(found_path)) embeddings.append(embedding) valid_entries.append(entry) valid_indices.append(i) # Check for multiple speakers if pipeline is active if diarization_pipeline is not None: # Diarize the file diar = diarization_pipeline(str(found_path)) speakers_found = diar.labels() if len(speakers_found) > 1: multi_speaker_files.append(str(found_path)) except Exception: pass if not embeddings: print("No embeddings extracted.") return embeddings = np.array(embeddings) print(f"Extracted {len(embeddings)} embeddings. Saving cache...") np.save(cache_path, embeddings) with open(cache_ids_path, 'w', encoding='utf-8') as f: json.dump(valid_indices, f) if diarization_pipeline is not None: multi_out = args.input.parent / "multi_speaker_detected.txt" with open(multi_out, "w", encoding="utf-8") as f: for mf in multi_speaker_files: f.write(mf + "\n") print(f"Found {len(multi_speaker_files)} files with multiple speakers. List saved to {multi_out}") print(f"Embeddings shape: {embeddings.shape}") # 4. Clustering (Scalable) # AgglomerativeClustering is O(N^2) memory, which crashes on 300k files (requires 300GB+ RAM). # We MUST use MiniBatchKMeans for this size. from sklearn.cluster import MiniBatchKMeans # Heuristic: Estimate number of speakers. n_clusters = args.num_clusters if n_clusters > len(embeddings): print(f"[WARNING] You asked for {n_clusters} clusters, but there are only {len(embeddings)} audio files!") n_clusters = len(embeddings) // 2 if n_clusters < 1: n_clusters = 1 print(f"Auto-adjusting clusters to: {n_clusters}") print(f"Clustering into {n_clusters} estimated groups using MiniBatchKMeans...") clustering = MiniBatchKMeans( n_clusters=n_clusters, batch_size=4096, random_state=42, init='random', n_init=1 ).fit(embeddings) labels = clustering.labels_ num_speakers = len(set(labels)) print(f"Found {num_speakers} cluster groups.") # 5. Save Result print(f"Saving to {args.output}...") args.output.parent.mkdir(parents=True, exist_ok=True) with args.output.open("w", encoding="utf-8") as f: for entry, label in zip(valid_entries, labels): entry['speaker'] = f"speaker_{label}" # Assign the new Cluster ID f.write(json.dumps(entry, ensure_ascii=False) + "\n") print("Done! You now have a properly labeled manifest.") if __name__ == "__main__": main()