handler.py

"""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)