"""Hugging Face ZeroGPU & Inference Endpoint Custom Handler. Combines faster-whisper (ASR) + pyannote 4.0 (Speaker Diarization) in a single call with WhisperX-style word/segment speaker alignment. Compatible with both Hugging Face Inference Endpoints and ZeroGPU Spaces via lazy model loading. """ import base64 import io import logging import os import tempfile from pathlib import Path from typing import Any, Dict, List, Optional, Tuple try: import numpy as np except ImportError: np = None try: import torch except ImportError: torch = None try: from faster_whisper import WhisperModel except ImportError: WhisperModel = None try: from pyannote.audio import Pipeline except ImportError: Pipeline = None try: from pydub import AudioSegment except ImportError: AudioSegment = None try: import soundfile as sf except ImportError: sf = None logger = logging.getLogger(__name__) logging.basicConfig(level=logging.INFO) class EndpointHandler: def __init__(self, path: str = "", lazy: bool = False): """Initializes the handler configuration. Args: path: Model directory path. lazy: If True, defers model loading until the first inference call inside the @spaces.GPU execution lifecycle. """ self.path = path self.lazy = lazy self.whisper_model = None self.diarization_pipeline = None # Model configuration from environment variables self.whisper_model_name = os.environ.get("WHISPER_MODEL", "large-v3") self.pyannote_model_name = os.environ.get( "PYANNOTE_MODEL", "pyannote/speaker-diarization-community-1" ) self.hf_token = ( os.getenv("HFTOKEN") or os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_TOKEN") or os.getenv("HF_API_TOKEN") ) if self.hf_token: os.environ["HF_TOKEN"] = self.hf_token os.environ["HUGGING_FACE_HUB_TOKEN"] = self.hf_token if not self.lazy: self._load_models() def _setup_cuda_env(self) -> None: """Configures library paths so CTranslate2 can dynamically find libcublas.so.12 and libcudnn.""" try: import nvidia.cublas.lib import nvidia.cudnn.lib cublas_lib_dir = os.path.dirname(nvidia.cublas.lib.__file__) cudnn_lib_dir = os.path.dirname(nvidia.cudnn.lib.__file__) current_ld = os.environ.get("LD_LIBRARY_PATH", "") new_dirs = [d for d in [cublas_lib_dir, cudnn_lib_dir] if d and os.path.isdir(d)] if new_dirs: os.environ["LD_LIBRARY_PATH"] = ":".join(new_dirs) + (f":{current_ld}" if current_ld else "") # Preload libcublas and libcublasLt into process address space import ctypes import glob for so_file in glob.glob(os.path.join(cublas_lib_dir, "libcublas*.so*")): try: ctypes.CDLL(so_file) except Exception: pass except Exception as e: logger.debug("CUDA runtime library path setup: %s", e) def _load_models(self) -> None: """Loads faster-whisper and pyannote models onto the active device (GPU or CPU).""" if self.whisper_model is not None and self.diarization_pipeline is not None: return self.device = "cuda" if (torch is not None and torch.cuda.is_available()) else "cpu" self.compute_type = "float16" if self.device == "cuda" else "int8" if self.device == "cuda": self._setup_cuda_env() logger.info( "Loading models into EndpointHandler on device=%s (compute_type=%s)...", self.device, self.compute_type, ) # 1. Load faster-whisper if self.whisper_model is None and WhisperModel is not None: logger.info("Loading faster-whisper model '%s' on %s...", self.whisper_model_name, self.device) self.whisper_model = WhisperModel( self.whisper_model_name, device=self.device, compute_type=self.compute_type, ) elif WhisperModel is None: logger.warning("faster-whisper is not installed.") # 2. Load pyannote.audio pipeline if self.diarization_pipeline is None and Pipeline is not None: logger.info("Loading pyannote diarization pipeline '%s'...", self.pyannote_model_name) try: self.diarization_pipeline = Pipeline.from_pretrained( self.pyannote_model_name, token=self.hf_token, ) if self.diarization_pipeline is not None and self.device == "cuda": self.diarization_pipeline.to(torch.device("cuda")) except Exception as e: logger.error( "Failed to authenticate/load pyannote pipeline '%s': %s", self.pyannote_model_name, e, ) raise RuntimeError( f"PyAnnote model '{self.pyannote_model_name}' failed to load: {str(e)}. " "Ensure HFTOKEN is set in Space Secrets and gated model terms are accepted." ) elif Pipeline is None: logger.warning("pyannote.audio is not installed.") logger.info("Model loading complete.") def _prepare_audio(self, data: Any) -> Tuple[str, float]: """Converts diverse payload types into a standard 16kHz mono WAV temporary file. Returns (temp_file_path, duration_seconds). """ raw_bytes: Optional[bytes] = None if isinstance(data, bytes): raw_bytes = data elif isinstance(data, dict): inputs = data.get("inputs") if isinstance(inputs, bytes): raw_bytes = inputs elif isinstance(inputs, str): if inputs.startswith("data:audio") or ";base64," in inputs: raw_bytes = base64.b64decode(inputs.split(";base64,")[-1]) elif len(inputs) > 500 and not inputs.startswith("http"): try: raw_bytes = base64.b64decode(inputs) except Exception: raw_bytes = None elif os.path.exists(inputs): with open(inputs, "rb") as f: raw_bytes = f.read() elif inputs.startswith("http://") or inputs.startswith("https://"): import urllib.request req = urllib.request.Request(inputs, headers={"User-Agent": "MeetPilot-HF-Handler/1.0"}) with urllib.request.urlopen(req) as resp: raw_bytes = resp.read() elif isinstance(data, str) and os.path.exists(data): with open(data, "rb") as f: raw_bytes = f.read() if raw_bytes is None: raise ValueError("No valid audio bytes or input file found in request payload.") # Convert to 16kHz mono WAV file temp_wav = tempfile.NamedTemporaryFile(suffix=".wav", delete=False) temp_path = temp_wav.name temp_wav.close() try: if AudioSegment is not None: audio = AudioSegment.from_file(io.BytesIO(raw_bytes)) audio = audio.set_frame_rate(16000).set_channels(1) audio.export(temp_path, format="wav") duration = len(audio) / 1000.0 return temp_path, duration elif sf is not None: audio_data, sr = sf.read(io.BytesIO(raw_bytes)) if len(audio_data.shape) > 1: audio_data = audio_data.mean(axis=1) sf.write(temp_path, audio_data, sr, format="WAV") duration = len(audio_data) / float(sr) return temp_path, duration else: with open(temp_path, "wb") as f: f.write(raw_bytes) return temp_path, 0.0 except Exception as e: logger.warning("Audio conversion fallback to raw write: %s", e) with open(temp_path, "wb") as f: f.write(raw_bytes) return temp_path, 0.0 def _diarize_audio(self, wav_path: str, min_speakers: Optional[int], max_speakers: Optional[int]) -> List[Dict[str, Any]]: """Runs pyannote diarization and returns chronological turns with start, end, speaker.""" if self.diarization_pipeline is None: logger.warning("Diarization pipeline not available; defaulting to single speaker.") return [] kwargs = {} if min_speakers is not None: try: kwargs["min_speakers"] = int(min_speakers) except (ValueError, TypeError): pass if max_speakers is not None: try: kwargs["max_speakers"] = int(max_speakers) except (ValueError, TypeError): pass try: # Provide in-memory waveform dictionary to pyannote to bypass torchcodec and avoid libnvrtc dependencies audio_input = wav_path if torch is not None and sf is not None: try: audio_data, sr = sf.read(wav_path, dtype="float32") if len(audio_data.shape) == 1: waveform = torch.from_numpy(audio_data).unsqueeze(0) # Shape: (1, samples) else: waveform = torch.from_numpy(audio_data.T) # Shape: (channels, samples) audio_input = {"waveform": waveform, "sample_rate": int(sr)} except Exception as load_exc: logger.debug("Soundfile tensor conversion fallback: %s", load_exc) audio_input = wav_path elif torch is not None: try: import torchaudio waveform, sr = torchaudio.load(wav_path) audio_input = {"waveform": waveform, "sample_rate": int(sr)} except Exception as load_exc: logger.debug("Torchaudio loading fallback: %s", load_exc) audio_input = wav_path diarization_output = self.diarization_pipeline(audio_input, **kwargs) turns: List[Dict[str, Any]] = [] # Support both PyAnnote 3.x and 4.x output formats if hasattr(diarization_output, "itertracks"): for turn, _, speaker in diarization_output.itertracks(yield_label=True): turns.append({ "start": float(turn.start), "end": float(turn.end), "speaker": str(speaker), }) elif hasattr(diarization_output, "speaker_diarization"): for turn, speaker in diarization_output.speaker_diarization: turns.append({ "start": float(turn.start), "end": float(turn.end), "speaker": str(speaker), }) else: for segment in diarization_output: turns.append({ "start": float(segment.start), "end": float(segment.end), "speaker": str(getattr(segment, "speaker", "SPEAKER_00")), }) return sorted(turns, key=lambda t: t["start"]) except Exception as exc: logger.error("Diarization failed: %s", exc) return [] def _align_words_with_diarization( self, whisper_segments: List[Any], diarization_turns: List[Dict[str, Any]], ) -> List[Dict[str, Any]]: """Aligns Whisper word-level timestamps with pyannote diarization turns using WhisperX overlap maximization.""" speaker_mapping: Dict[str, str] = {} speaker_counter = 1 def get_clean_speaker_name(raw_spk: str) -> str: nonlocal speaker_counter if not raw_spk: return "Speaker 1" if raw_spk not in speaker_mapping: speaker_mapping[raw_spk] = f"Speaker {speaker_counter}" speaker_counter += 1 return speaker_mapping[raw_spk] all_words: List[Dict[str, Any]] = [] for seg in whisper_segments: seg_words = getattr(seg, "words", None) if seg_words: for w in seg_words: word_text = getattr(w, "word", "") w_start = getattr(w, "start", seg.start) w_end = getattr(w, "end", seg.end) if word_text.strip(): all_words.append({ "word": word_text, "start": float(w_start), "end": float(w_end), "seg_start": float(seg.start), "seg_end": float(seg.end), }) else: seg_text = getattr(seg, "text", "").strip() if seg_text: all_words.append({ "word": seg_text, "start": float(seg.start), "end": float(seg.end), "seg_start": float(seg.start), "seg_end": float(seg.end), }) if not all_words: return [] if not diarization_turns: return [ { "speaker": "Speaker 1", "start_time": round(float(seg.start), 3), "end_time": round(float(seg.end), 3), "text": seg.text.strip(), } for seg in whisper_segments if getattr(seg, "text", "").strip() ] last_known_speaker = "Speaker 1" for w in all_words: w_start = w["start"] w_end = w["end"] best_speaker = None max_overlap = 0.0 for turn in diarization_turns: t_start = turn["start"] t_end = turn["end"] overlap = max(0.0, min(w_end, t_end) - max(w_start, t_start)) if overlap > max_overlap: max_overlap = overlap best_speaker = turn["speaker"] if not best_speaker or max_overlap <= 0.0: mid_point = (w_start + w_end) / 2.0 closest_turn = min( diarization_turns, key=lambda t: min(abs(mid_point - t["start"]), abs(mid_point - t["end"])), ) dist = min(abs(mid_point - closest_turn["start"]), abs(mid_point - closest_turn["end"])) if dist <= 1.5: best_speaker = closest_turn["speaker"] else: best_speaker = last_known_speaker clean_spk = get_clean_speaker_name(best_speaker) w["speaker"] = clean_spk last_known_speaker = clean_spk final_segments: List[Dict[str, Any]] = [] current_speaker = all_words[0]["speaker"] current_start = all_words[0]["start"] current_end = all_words[0]["end"] current_words: List[str] = [all_words[0]["word"]] for w in all_words[1:]: spk = w["speaker"] w_start = w["start"] w_end = w["end"] w_text = w["word"] is_same_speaker = (spk == current_speaker) gap = max(0.0, w_start - current_end) if is_same_speaker and gap <= 2.0: current_words.append(w_text) current_end = max(current_end, w_end) else: seg_text = "".join(current_words).strip() if seg_text: final_segments.append({ "speaker": current_speaker, "start_time": round(current_start, 3), "end_time": round(current_end, 3), "text": seg_text, }) current_speaker = spk current_start = w_start current_end = w_end current_words = [w_text] if current_words: seg_text = "".join(current_words).strip() if seg_text: final_segments.append({ "speaker": current_speaker, "start_time": round(current_start, 3), "end_time": round(current_end, 3), "text": seg_text, }) return final_segments def __call__(self, data: Any) -> Dict[str, Any]: """Main inference entrypoint. Accepts audio payload, ensures models are loaded on the active device, runs faster-whisper + pyannote diarization, and returns structured JSON. """ # Ensure models are loaded (especially on ZeroGPU when GPU is granted) self._load_models() temp_wav_path = None try: # 1. Parse optional parameters parameters = {} if isinstance(data, dict) and "parameters" in data: parameters = data.get("parameters") or {} language = parameters.get("language") min_speakers = parameters.get("min_speakers") max_speakers = parameters.get("max_speakers") initial_prompt = parameters.get("initial_prompt") # 2. Extract and standardize audio to 16kHz WAV temp_wav_path, duration = self._prepare_audio(data) # 3. Step A: Faster-Whisper ASR with word timestamps if self.whisper_model is None: raise RuntimeError("Whisper model is not initialized.") logger.info("Running faster-whisper transcription...") whisper_segments_gen, info = self.whisper_model.transcribe( temp_wav_path, beam_size=5, word_timestamps=True, language=language, initial_prompt=initial_prompt, vad_filter=True, ) whisper_segments = list(whisper_segments_gen) logger.info( "Whisper transcribed %d segments (language=%s, duration=%.1fs).", len(whisper_segments), getattr(info, "language", "unknown"), getattr(info, "duration", duration), ) # 4. Step B: PyAnnote Diarization logger.info("Running pyannote diarization...") diarization_turns = self._diarize_audio( temp_wav_path, min_speakers=min_speakers, max_speakers=max_speakers, ) logger.info("PyAnnote extracted %d diarization turns.", len(diarization_turns)) # 5. Step C: WhisperX-style timestamp alignment final_segments = self._align_words_with_diarization( whisper_segments=whisper_segments, diarization_turns=diarization_turns, ) return { "segments": final_segments, "language": getattr(info, "language", "en"), "duration": round(getattr(info, "duration", duration), 2), } except Exception as exc: logger.exception("Error processing audio in EndpointHandler: %s", exc) raise exc finally: if temp_wav_path and os.path.exists(temp_wav_path): try: os.unlink(temp_wav_path) except Exception as clean_err: logger.warning("Failed to clean up temp file %s: %s", temp_wav_path, clean_err)