""" Voice Extractor CPU - Pipeline module Identifies, isolates, and transcribes a target speaker from multi-speaker audio. Based on https://github.com/ReisCook/Voice_Extractor Heavy imports (torch, torchaudio, pyannote, speechbrain, whisper) are LAZY — imported inside functions only. This keeps startup fast when only the audiosplitter tab is used. """ import os import sys import shutil import tempfile import logging import re import csv import time from pathlib import Path import numpy as np # --------------------------------------------------------------------------- # Force CPU # --------------------------------------------------------------------------- os.environ["CUDA_VISIBLE_DEVICES"] = "" # --------------------------------------------------------------------------- # Logging # --------------------------------------------------------------------------- logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") log = logging.getLogger("voice_extractor_cpu") # --------------------------------------------------------------------------- # Model paths (bundled locally — no HF_TOKEN needed) # --------------------------------------------------------------------------- _APP_DIR = os.path.dirname(os.path.abspath(__file__)) PYANNOTE_DIR = os.path.join(_APP_DIR, "models", "audiosplitter", "pyannote") BANDIT_ONNX_PATH = os.path.join(_APP_DIR, "models", "voice_extractor", "bandit_v2_speech_fp32.onnx") HF_TOKEN = os.environ.get("HF_TOKEN") # fallback, not required # --------------------------------------------------------------------------- # Intercept hf_hub_download to serve bundled PyAnnote models locally. # Also fixes pyannote's use_auth_token → token compat issue. # --------------------------------------------------------------------------- import huggingface_hub as _hfhub _orig_hf_hub_download = _hfhub.hf_hub_download # Map HF repo IDs to local model directories _LOCAL_MODEL_MAP = { "pyannote/speaker-diarization-3.1": os.path.join(PYANNOTE_DIR, "speaker-diarization-3.1"), "pyannote/segmentation-3.0": os.path.join(PYANNOTE_DIR, "segmentation-3.0"), "pyannote/wespeaker-voxceleb-resnet34-LM": os.path.join(PYANNOTE_DIR, "wespeaker-voxceleb-resnet34-LM"), "pyannote/segmentation": os.path.join(PYANNOTE_DIR, "segmentation"), "pyannote/overlapped-speech-detection": os.path.join(PYANNOTE_DIR, "overlapped-speech-detection"), } def _patched_hf_hub_download(*args, **kwargs): # Fix use_auth_token → token if "use_auth_token" in kwargs: kwargs["token"] = kwargs.pop("use_auth_token") # Intercept PyAnnote model downloads → serve from local repo_id = args[0] if args else kwargs.get("repo_id", "") filename = args[1] if len(args) > 1 else kwargs.get("filename", "") if repo_id in _LOCAL_MODEL_MAP: local_path = os.path.join(_LOCAL_MODEL_MAP[repo_id], filename) if os.path.exists(local_path): log.info(f"[LOCAL] {repo_id}/{filename} → {local_path}") return local_path else: log.warning(f"[LOCAL] {repo_id}/{filename} not found locally, falling back to HF") return _orig_hf_hub_download(*args, **kwargs) _hfhub.hf_hub_download = _patched_hf_hub_download # Also patch cached_download if it exists (used by some older pyannote internals) if hasattr(_hfhub, "cached_download"): _orig_cached_download = _hfhub.cached_download def _patched_cached_download(*args, **kwargs): if "use_auth_token" in kwargs: kwargs["token"] = kwargs.pop("use_auth_token") return _orig_cached_download(*args, **kwargs) _hfhub.cached_download = _patched_cached_download # --------------------------------------------------------------------------- # Lazy-loaded models (singleton pattern) # --------------------------------------------------------------------------- _models = { "wespeaker_rvector": None, "wespeaker_gemini": None, "speechbrain": None, "whisper": None, } def _get_wespeaker_models(): """Load WeSpeaker models (cached).""" if _models["wespeaker_rvector"] is None: import wespeaker log.info("Loading WeSpeaker english model (r-vector)...") m = wespeaker.load_model("english") m.set_device("cpu") _models["wespeaker_rvector"] = m # Use same model for gemini slot on CPU to save memory _models["wespeaker_gemini"] = m log.info("WeSpeaker model loaded.") return {"rvector": _models["wespeaker_rvector"], "gemini": _models["wespeaker_gemini"]} def _get_speechbrain_model(): """Load SpeechBrain ECAPA-TDNN (cached).""" if _models["speechbrain"] is None: try: from speechbrain.inference.speaker import SpeakerRecognition log.info("Loading SpeechBrain ECAPA-TDNN...") cache_dir = Path(tempfile.gettempdir()) / "voice_extractor_sb_cache" / "spkrec-ecapa-voxceleb" cache_dir.mkdir(parents=True, exist_ok=True) model = SpeakerRecognition.from_hparams( source="speechbrain/spkrec-ecapa-voxceleb", savedir=str(cache_dir), run_opts={"device": "cpu"}, ) model.eval() _models["speechbrain"] = model log.info("SpeechBrain ECAPA-TDNN loaded.") except Exception as e: log.warning(f"SpeechBrain load failed: {e}. Verification will use WeSpeaker only.") _models["speechbrain"] = False # sentinel: tried and failed if _models["speechbrain"] is False: return None return _models["speechbrain"] def _get_whisper_model(model_name: str = "base"): """Load Whisper model (cached).""" if _models["whisper"] is None: import whisper log.info(f"Loading Whisper '{model_name}' on CPU...") _models["whisper"] = whisper.load_model(model_name, device="cpu") log.info("Whisper model loaded.") return _models["whisper"] # --------------------------------------------------------------------------- # Bandit-v2 vocal separation (ONNX, speech stem only) # --------------------------------------------------------------------------- _bandit_session = None def _get_bandit_session(): """Load Bandit-v2 ONNX session (cached).""" global _bandit_session if _bandit_session is None: if not os.path.exists(BANDIT_ONNX_PATH): log.warning(f"Bandit-v2 ONNX not found at {BANDIT_ONNX_PATH}") return None import onnxruntime as ort log.info("Loading Bandit-v2 ONNX model...") _bandit_session = ort.InferenceSession(BANDIT_ONNX_PATH, providers=["CPUExecutionProvider"]) log.info("Bandit-v2 ONNX loaded.") return _bandit_session def bandit_vocal_separation(input_path: Path, output_path: Path, chunk_seconds: float = 8.0, fs: int = 48000, progress_cb=None) -> Path | None: """Separate vocals from music/SFX using Bandit-v2 ONNX. STFT/iSTFT in PyTorch, core neural net in ONNX Runtime. Returns path to separated speech WAV, or None on failure. """ import torch import librosa import soundfile as sf sess = _get_bandit_session() if sess is None: return None import torchaudio.transforms as T if progress_cb: progress_cb(0.02, "Loading audio for vocal separation...") # Load and resample to 48kHz mono y, sr = librosa.load(str(input_path), sr=fs, mono=True) audio = torch.from_numpy(y).unsqueeze(0) # (1, samples) n_samples = audio.shape[-1] log.info(f"Bandit-v2: {n_samples/fs:.1f}s audio @ {fs}Hz") # STFT params (must match training config) n_fft, hop_length, win_length = 2048, 512, 2048 stft = T.Spectrogram(n_fft=n_fft, win_length=win_length, hop_length=hop_length, power=None, normalized=True, center=True, pad_mode="reflect") istft = T.InverseSpectrogram(n_fft=n_fft, win_length=win_length, hop_length=hop_length, normalized=True, center=True) # Chunked overlap-add inference chunk_samples = int(chunk_seconds * fs) hop_samples = fs # 1s hop overlap = chunk_samples - hop_samples front_pad = 2 * overlap window = torch.hann_window(chunk_samples) / (chunk_samples / (2 * hop_samples)) padded = torch.nn.functional.pad(audio, (front_pad, front_pad), mode="constant") total_len = padded.shape[-1] starts = list(range(0, total_len - chunk_samples + 1, hop_samples)) output_buf = torch.zeros_like(padded) log.info(f"Bandit-v2: {len(starts)} chunks, {chunk_seconds}s each") for idx, start in enumerate(starts): if progress_cb and idx % max(1, len(starts)//10) == 0: pct = 0.02 + 0.15 * (idx / len(starts)) progress_cb(pct, f"Vocal separation: chunk {idx+1}/{len(starts)}...") chunk = padded[:, start:start+chunk_samples] # (1, chunk_samples) # STFT → complex spectrogram spec = stft(chunk) # (1, freq, time) complex spec_real = spec.real.unsqueeze(0).numpy() # (1, 1, freq, time) spec_imag = spec.imag.unsqueeze(0).numpy() # ONNX core inference speech_real, speech_imag = sess.run(None, { "spec_real": spec_real.astype(np.float32), "spec_imag": spec_imag.astype(np.float32), }) # iSTFT → waveform masked_spec = torch.complex( torch.from_numpy(speech_real[0]), # (1, freq, time) torch.from_numpy(speech_imag[0]), ) chunk_out = istft(masked_spec, chunk_samples) # (1, chunk_samples) output_buf[:, start:start+chunk_samples] += chunk_out * window # Trim padding speech_audio = output_buf[:, front_pad:front_pad+n_samples] # Save sf.write(str(output_path), speech_audio.squeeze().numpy(), fs) log.info(f"Bandit-v2: saved {output_path.name}") if progress_cb: progress_cb(0.18, "Vocal separation complete.") return output_path # --------------------------------------------------------------------------- # Audio utilities # --------------------------------------------------------------------------- def ff_convert(src: Path, dst: Path, sr: int = 16000, ac: int = 1, start: float | None = None, end: float | None = None): """Convert/trim audio with ffmpeg.""" import ffmpeg inp_kwargs = {} if start is not None: inp_kwargs["ss"] = start if end is not None: inp_kwargs["to"] = end ( ffmpeg.input(str(src), **inp_kwargs) .output(str(dst), acodec="pcm_s16le", ac=ac, ar=sr) .overwrite_output() .run(quiet=True, capture_stdout=True, capture_stderr=True) ) def cosine_sim(a, b) -> float: na, nb = np.linalg.norm(a), np.linalg.norm(b) if na == 0 or nb == 0: return 0.0 return float(np.dot(a, b) / (na * nb)) # --------------------------------------------------------------------------- # Pipeline stages # --------------------------------------------------------------------------- def prepare_reference(ref_path: Path, tmp_dir: Path) -> Path: """Convert reference to 16 kHz mono WAV.""" out = tmp_dir / "reference_16k.wav" ff_convert(ref_path, out, sr=16000, ac=1) return out def diarize(audio_path: Path, tmp_dir: Path, dry_run_sec: int | None = None, progress_cb=None) -> "Annotation | None": """Run PyAnnote speaker diarization 3.1 (local models, no HF_TOKEN needed).""" import torch from pyannote.audio import Pipeline as PyannotePipeline DEVICE = torch.device("cpu") if progress_cb: progress_cb(0.10, "Loading diarization model...") # hf_hub_download interceptor serves local models automatically pipeline = PyannotePipeline.from_pretrained( "pyannote/speaker-diarization-3.1", use_auth_token=HF_TOKEN or True ) # Force CPU pipeline = pipeline.to(DEVICE) target = audio_path if dry_run_sec: cut = tmp_dir / "diar_cut.wav" ff_convert(audio_path, cut, sr=16000, ac=1, end=dry_run_sec) target = cut if progress_cb: progress_cb(0.15, "Diarizing speakers...") result = pipeline({"uri": target.stem, "audio": str(target)}) n = len(result.labels()) log.info(f"Diarization found {n} speakers.") return result def detect_overlaps(audio_path: Path, tmp_dir: Path, dry_run_sec: int | None = None, progress_cb=None) -> "Timeline": """Run PyAnnote overlapped speech detection.""" import torch from pyannote.audio import Pipeline as PyannotePipeline from pyannote.core import Timeline DEVICE = torch.device("cpu") if progress_cb: progress_cb(0.30, "Loading overlap detection model...") try: # hf_hub_download interceptor serves local models automatically osd = PyannotePipeline.from_pretrained( "pyannote/overlapped-speech-detection", use_auth_token=HF_TOKEN or True ) osd = osd.to(DEVICE) except Exception as e: log.warning(f"OSD load failed: {e}. Proceeding without overlap detection.") return Timeline() target = audio_path if dry_run_sec: cut = tmp_dir / "osd_cut.wav" ff_convert(audio_path, cut, sr=16000, ac=1, end=dry_run_sec) target = cut if progress_cb: progress_cb(0.35, "Detecting overlaps...") try: result = osd({"uri": target.stem, "audio": str(target)}) if isinstance(result, Timeline): return result.support() # If Annotation, extract overlap label from pyannote.core import Annotation if isinstance(result, Annotation): tl = Timeline() if "overlap" in result.labels(): tl.update(result.label_timeline("overlap")) return tl.support() return Timeline() except Exception as e: log.warning(f"OSD failed: {e}") return Timeline() def identify_target(diar_annotation, audio_path: Path, ref_16k: Path, target_name: str, progress_cb=None) -> str | None: """Identify which diarized speaker matches the reference.""" import ffmpeg ws_models = _get_wespeaker_models() ws = ws_models["rvector"] if ws is None: return None if progress_cb: progress_cb(0.45, "Identifying target speaker...") ref_emb = ws.extract_embedding(str(ref_16k)) labels = diar_annotation.labels() if not labels: return None best_label, best_score = None, -1.0 with tempfile.TemporaryDirectory(prefix="spk_id_") as td: td_path = Path(td) for label in labels: tl = diar_annotation.label_timeline(label) if not tl: continue # Gather up to 20s of audio for this speaker segs = [] total = 0.0 for seg in tl: if total >= 20.0: break seg_path = td_path / f"{label}_{len(segs)}.wav" try: ff_convert(audio_path, seg_path, sr=16000, ac=1, start=seg.start, end=seg.end) if seg_path.exists() and seg_path.stat().st_size > 0: segs.append(seg_path) total += seg.duration except Exception: continue if not segs: continue # Concatenate if multiple segments if len(segs) == 1: concat_path = segs[0] else: concat_path = td_path / f"{label}_concat.wav" list_file = td_path / f"{label}_list.txt" list_file.write_text( "\n".join(f"file '{p.resolve().as_posix()}'" for p in segs) ) try: ( ffmpeg.input(str(list_file), format="concat", safe=0) .output(str(concat_path), acodec="pcm_s16le", ar=16000, ac=1) .overwrite_output() .run(quiet=True, capture_stdout=True, capture_stderr=True) ) except Exception: concat_path = segs[0] try: spk_emb = ws.extract_embedding(str(concat_path)) sim = cosine_sim(ref_emb, spk_emb) log.info(f" Speaker {label}: similarity = {sim:.4f}") if sim > best_score: best_score = sim best_label = label except Exception as e: log.warning(f"Embedding failed for {label}: {e}") if best_label: log.info(f"Identified '{target_name}' as {best_label} (score {best_score:.4f})") return best_label def verify_segment(seg_path: Path, ref_path: Path, use_speechbrain: bool = True) -> tuple[float, dict]: """Multi-model speaker verification on a single segment.""" import torch import librosa ws_models = _get_wespeaker_models() scores = {"wespeaker": 0.0, "speechbrain": 0.0} # WeSpeaker try: ws = ws_models["rvector"] ref_emb = ws.extract_embedding(str(ref_path)) seg_emb = ws.extract_embedding(str(seg_path)) scores["wespeaker"] = cosine_sim(ref_emb, seg_emb) except Exception as e: log.debug(f"WeSpeaker verify fail: {e}") # SpeechBrain if use_speechbrain: sb = _get_speechbrain_model() if sb is not None: try: score_t, _ = sb.verify_files( str(ref_path.resolve()).replace("\\", "/"), str(seg_path.resolve()).replace("\\", "/"), ) scores["speechbrain"] = score_t.item() except Exception as e: log.debug(f"SpeechBrain verify fail: {e}") # VAD check vad_factor = 1.0 try: y, sr = librosa.load(seg_path, sr=16000, mono=True) if len(y) > 0: vad_model, utils = torch.hub.load( "snakers4/silero-vad", "silero_vad", force_reload=False, trust_repo=True, verbose=False, onnx=False, ) get_ts = utils[0] audio_t = torch.FloatTensor(y) ts = get_ts(audio_t, vad_model, sampling_rate=16000, threshold=0.5) speech_dur = sum(d["end"] - d["start"] for d in ts) / 16000 total_dur = len(y) / 16000 ratio = speech_dur / total_dur if total_dur > 0 else 0 vad_factor = 1.0 if ratio >= 0.6 else 0.1 except Exception: pass # Combine if scores["speechbrain"] > 0: combined = (scores["wespeaker"] * 0.5 + scores["speechbrain"] * 0.5) * vad_factor else: combined = scores["wespeaker"] * vad_factor return combined, scores def extract_and_verify( diar_annotation, target_label: str, overlap_tl, audio_path: Path, ref_16k: Path, target_name: str, output_dir: Path, tmp_dir: Path, threshold: float = 0.7, min_duration: float = 1.0, merge_gap: float = 0.25, output_sr: int = 44100, use_speechbrain: bool = True, progress_cb=None, ) -> tuple[list[Path], list[Path]]: """Slice target solo segments, verify identity, return verified/rejected paths.""" from pyannote.core import Segment, Timeline # Get target solo timeline (minus overlaps) target_tl = diar_annotation.label_timeline(target_label).support() solo_tl = target_tl.extrude(overlap_tl.support()) if overlap_tl else target_tl # Merge nearby segments segs = sorted(list(solo_tl), key=lambda s: s.start) merged = [] if segs: cur = segs[0] for nxt in segs[1:]: if nxt.start <= cur.end + merge_gap and nxt.end > cur.end: cur = Segment(cur.start, nxt.end) elif nxt.start > cur.end + merge_gap: merged.append(cur) cur = nxt merged.append(cur) # Duration filter merged = [s for s in merged if s.duration >= min_duration] if not merged: return [], [] verified_dir = output_dir / "verified" rejected_dir = output_dir / "rejected" verified_dir.mkdir(parents=True, exist_ok=True) rejected_dir.mkdir(parents=True, exist_ok=True) verif_tmp = tmp_dir / "verif_16k" hq_tmp = tmp_dir / "hq" verif_tmp.mkdir(parents=True, exist_ok=True) hq_tmp.mkdir(parents=True, exist_ok=True) verified_paths = [] rejected_paths = [] total = len(merged) for i, seg in enumerate(merged): if progress_cb: pct = 0.55 + 0.25 * (i / max(total, 1)) progress_cb(pct, f"Verifying segment {i+1}/{total}...") s_str = f"{seg.start:.3f}".replace(".", "p") e_str = f"{seg.end:.3f}".replace(".", "p") base = f"{i:04d}_{s_str}s_to_{e_str}s" seg_16k = verif_tmp / f"{base}.wav" seg_hq = hq_tmp / f"{base}_hq.wav" try: ff_convert(audio_path, seg_16k, sr=16000, ac=1, start=seg.start, end=seg.end) ff_convert(audio_path, seg_hq, sr=output_sr, ac=1, start=seg.start, end=seg.end) except Exception as e: log.warning(f"Slice failed for segment {i}: {e}") continue if not seg_16k.exists() or seg_16k.stat().st_size == 0: continue score, _ = verify_segment(seg_16k, ref_16k, use_speechbrain=use_speechbrain) safe_name = re.sub(r'[<>:"/\\|?*]', "", target_name).replace(" ", "_") if score >= threshold: dst = verified_dir / f"{safe_name}_verified_{base}.wav" shutil.copy(seg_hq, dst) verified_paths.append(dst) else: dst = rejected_dir / f"{safe_name}_rejected_{base}_score{score:.3f}.wav" shutil.copy(seg_hq, dst) rejected_paths.append(dst) # Cleanup temp seg_16k.unlink(missing_ok=True) seg_hq.unlink(missing_ok=True) log.info(f"Verified: {len(verified_paths)}, Rejected: {len(rejected_paths)}") return verified_paths, rejected_paths def transcribe_files( segment_paths: list[Path], whisper_model_name: str = "base", language: str = "en", progress_cb=None, ) -> list[dict]: """Transcribe audio segments with Whisper. Returns list of dicts.""" import librosa if not segment_paths: return [] model = _get_whisper_model(whisper_model_name) results = [] total = len(segment_paths) for i, p in enumerate(segment_paths): if progress_cb: pct = 0.82 + 0.15 * (i / max(total, 1)) progress_cb(pct, f"Transcribing {i+1}/{total}...") if not p.exists() or p.stat().st_size == 0: continue try: r = model.transcribe(str(p), fp16=False, language=language if language != "auto" else None) text = r["text"].strip() except Exception as e: text = f"[error: {e}]" log.warning(f"Transcribe fail for {p.name}: {e}") dur = librosa.get_duration(path=p) results.append({"file": p.name, "duration_s": round(dur, 2), "text": text}) return results def concatenate_verified(paths: list[Path], output_path: Path, silence_s: float = 0.25, sr: int = 44100): """Concatenate verified segments with silence gaps.""" import ffmpeg import soundfile as sf if not paths: return None # Sort by segment start time from filename def sort_key(p): m = re.search(r"(\d+p\d+)s_to_", p.name) if m: return float(m.group(1).replace("p", ".")) return 0.0 paths = sorted(paths, key=sort_key) with tempfile.TemporaryDirectory(prefix="concat_") as td: td_path = Path(td) # Create silence file silence_file = td_path / "silence.wav" if silence_s > 0: ( ffmpeg.input(f"anullsrc=channel_layout=mono:sample_rate={sr}", format="lavfi", t=str(silence_s)) .output(str(silence_file), acodec="pcm_s16le", ar=sr, ac=1) .overwrite_output() .run(quiet=True, capture_stdout=True, capture_stderr=True) ) # Build concat list lines = [] for i, p in enumerate(paths): if i > 0 and silence_s > 0 and silence_file.exists(): lines.append(f"file '{silence_file.resolve().as_posix()}'") lines.append(f"file '{p.resolve().as_posix()}'") list_file = td_path / "list.txt" list_file.write_text("\n".join(lines)) try: ( ffmpeg.input(str(list_file), format="concat", safe=0) .output(str(output_path), acodec="pcm_s16le", ar=sr, ac=1) .overwrite_output() .run(quiet=True, capture_stdout=True, capture_stderr=True) ) return output_path except Exception as e: log.error(f"Concat failed: {e}") return None # --------------------------------------------------------------------------- # Main pipeline (entry point for app.py) # --------------------------------------------------------------------------- def run_voice_extractor_pipeline( input_audio_path: str, reference_audio_path: str, target_name: str, whisper_model: str = "base", language: str = "en", verification_threshold: float = 0.7, min_duration: float = 1.0, merge_gap: float = 0.25, output_sr: int = 44100, use_speechbrain: bool = True, dry_run: bool = False, use_bandit: bool = False, progress=None, ): """Full voice extraction pipeline. Returns (status, audio_preview, output_files, transcript_text).""" import gradio as gr if not input_audio_path: raise gr.Error("Please upload an input audio file.") if not reference_audio_path: raise gr.Error("Please upload a reference audio clip of the target speaker.") if not target_name or not target_name.strip(): raise gr.Error("Please enter a target speaker name.") start_time = time.time() def progress_cb(pct, msg): try: if progress is not None: progress(pct, desc=msg) except Exception: pass progress_cb(0.02, "Setting up...") input_p = Path(input_audio_path) ref_p = Path(reference_audio_path) work_dir = Path(tempfile.mkdtemp(prefix="voice_ext_")) tmp_dir = work_dir / "tmp" output_dir = work_dir / "output" tmp_dir.mkdir(parents=True, exist_ok=True) output_dir.mkdir(parents=True, exist_ok=True) dry_sec = 60 if dry_run else None try: # Stage 0a: Trim audio for dry run BEFORE any processing actual_input = input_p if dry_sec: trimmed = tmp_dir / f"{input_p.stem}_trimmed.wav" ff_convert(input_p, trimmed, sr=16000, ac=1, end=dry_sec) actual_input = trimmed log.info(f"Dry run: trimmed to {dry_sec}s") # Stage 0b (optional): Bandit-v2 vocal separation source_audio = actual_input if use_bandit: progress_cb(0.02, "Starting vocal separation (Bandit-v2, slow on CPU)...") vocals_path = tmp_dir / f"{input_p.stem}_vocals.wav" result = bandit_vocal_separation(actual_input, vocals_path, progress_cb=progress_cb) if result and result.exists(): source_audio = result log.info(f"Using Bandit-v2 vocals for downstream: {result.name}") else: log.warning("Bandit-v2 failed, using original audio.") # Stage 1: Prepare reference progress_cb(0.05, "Preparing reference audio...") ref_16k = prepare_reference(ref_p, tmp_dir) # Stage 2: Diarization (audio already trimmed if dry_run) progress_cb(0.08, "Starting diarization...") diar = diarize(source_audio, tmp_dir, progress_cb=progress_cb) if diar is None or not diar.labels(): raise gr.Error("Diarization found no speakers. Check your audio file.") # Stage 3: Overlap detection progress_cb(0.28, "Detecting overlapping speech...") overlap_tl = detect_overlaps(source_audio, tmp_dir, progress_cb=progress_cb) # Stage 4: Identify target progress_cb(0.42, "Identifying target speaker...") target_label = identify_target(diar, source_audio, ref_16k, target_name, progress_cb=progress_cb) if not target_label: raise gr.Error( f"Could not identify '{target_name}' among diarized speakers. " "Try a cleaner/longer reference clip." ) # Stage 5: Extract & verify progress_cb(0.50, "Extracting and verifying segments...") verified, rejected = extract_and_verify( diar, target_label, overlap_tl, source_audio, ref_16k, target_name, output_dir, tmp_dir, threshold=verification_threshold, min_duration=min_duration, merge_gap=merge_gap, output_sr=output_sr, use_speechbrain=use_speechbrain, progress_cb=progress_cb, ) if not verified and not rejected: raise gr.Error("No speech segments found for the target speaker.") # Stage 6: Transcribe verified segments progress_cb(0.80, "Transcribing verified segments...") transcripts = transcribe_files( verified, whisper_model_name=whisper_model, language=language, progress_cb=progress_cb, ) # Stage 7: Concatenate verified concat_path = None if verified: progress_cb(0.96, "Concatenating verified segments...") concat_path = output_dir / "concatenated_verified.wav" concatenate_verified(verified, concat_path, silence_s=0.25, sr=output_sr) # Build outputs progress_cb(0.98, "Packaging results...") # Collect all output files for download output_files = [] if concat_path and concat_path.exists(): output_files.append(str(concat_path)) for p in sorted(verified): output_files.append(str(p)) # Transcript text transcript_lines = [] for t in transcripts: transcript_lines.append(f"[{t['file']}] ({t['duration_s']}s): {t['text']}") transcript_text = "\n\n".join(transcript_lines) if transcript_lines else "No transcripts generated." # Save transcript CSV if transcripts: csv_path = output_dir / "transcripts.csv" with csv_path.open("w", newline="", encoding="utf-8") as f: w = csv.writer(f) w.writerow(["filename", "duration_s", "transcript"]) for t in transcripts: w.writerow([t["file"], t["duration_s"], t["text"]]) output_files.append(str(csv_path)) elapsed = time.time() - start_time status = ( f"Done in {elapsed:.0f}s. " f"Found {len(diar.labels())} speakers. " f"Identified '{target_name}' as {target_label}. " f"Verified: {len(verified)} segments, Rejected: {len(rejected)} segments." ) # Audio preview = concatenated file if available, else first verified segment audio_preview = None if concat_path and concat_path.exists(): audio_preview = str(concat_path) elif verified: audio_preview = str(verified[0]) progress_cb(1.0, "Complete!") return status, audio_preview, output_files, transcript_text except gr.Error: raise except Exception as e: log.exception("Pipeline error") raise gr.Error(f"Pipeline error: {e}")