voice-to-notes/python/voice_to_notes/services/diarize.py

"""Diarization service — pyannote.audio speaker identification."""

from __future__ import annotations

import os
import subprocess
import sys
import tempfile
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

# Disable pyannote telemetry — it has a bug in v4.0.4 where
# np.isfinite(None) crashes when max_speakers is not set.
os.environ.setdefault("PYANNOTE_METRICS_ENABLED", "false")

from voice_to_notes.ipc.messages import progress_message
from voice_to_notes.ipc.protocol import write_message


def _ensure_wav(file_path: str) -> tuple[str, str | None]:
    """Convert audio to 16kHz mono WAV if needed.

    pyannote.audio v4.0.4 has a bug where its AudioDecoder returns
    duration=None for some formats (FLAC, etc.), causing crashes.
    Converting to WAV ensures the duration header is always present.

    Returns:
        (path_to_use, temp_path_or_None)
        If conversion was needed, temp_path is the WAV file to clean up.
    """
    ext = Path(file_path).suffix.lower()
    if ext == ".wav":
        return file_path, None

    tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    tmp.close()
    try:
        subprocess.run(
            [
                "ffmpeg", "-y", "-i", file_path,
                "-ar", "16000", "-ac", "1", "-c:a", "pcm_s16le",
                tmp.name,
            ],
            check=True,
            capture_output=True,
        )
        print(
            f"[sidecar] Converted {ext} to WAV for diarization",
            file=sys.stderr,
            flush=True,
        )
        return tmp.name, tmp.name
    except (subprocess.CalledProcessError, FileNotFoundError) as e:
        # ffmpeg not available or failed — try original file and hope for the best
        print(
            f"[sidecar] WAV conversion failed ({e}), using original file",
            file=sys.stderr,
            flush=True,
        )
        os.unlink(tmp.name)
        return file_path, None


@dataclass
class SpeakerSegment:
    """A time span assigned to a speaker."""

    speaker: str
    start_ms: int
    end_ms: int


@dataclass
class DiarizationResult:
    """Full diarization output."""

    speaker_segments: list[SpeakerSegment] = field(default_factory=list)
    num_speakers: int = 0
    speakers: list[str] = field(default_factory=list)


class DiarizeService:
    """Handles speaker diarization via pyannote.audio."""

    def __init__(self) -> None:
        self._pipeline: Any = None

    def _ensure_pipeline(self, hf_token: str | None = None) -> Any:
        """Load the pyannote diarization pipeline (lazy)."""
        if self._pipeline is not None:
            return self._pipeline

        print("[sidecar] Loading pyannote diarization pipeline...", file=sys.stderr, flush=True)

        # Use token from argument, fall back to environment variable
        if not hf_token:
            hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") or None

        # Persist token globally so ALL huggingface_hub sub-downloads use auth.
        # Pyannote has internal dependencies that don't forward the token= param.
        if hf_token:
            os.environ["HF_TOKEN"] = hf_token
            import huggingface_hub
            huggingface_hub.login(token=hf_token, add_to_git_credential=False)

        models = [
            "pyannote/speaker-diarization-3.1",
            "pyannote/speaker-diarization",
        ]

        last_error: Exception | None = None
        for model_name in models:
            try:
                from pyannote.audio import Pipeline

                self._pipeline = Pipeline.from_pretrained(model_name, token=hf_token)
                print(f"[sidecar] Loaded diarization model: {model_name}", file=sys.stderr, flush=True)
                return self._pipeline
            except Exception as e:
                last_error = e
                print(
                    f"[sidecar] Warning: Could not load {model_name}: {e}",
                    file=sys.stderr,
                    flush=True,
                )

        raise RuntimeError(
            "pyannote.audio pipeline not available. "
            "You may need to accept the model license at "
            "https://huggingface.co/pyannote/speaker-diarization-3.1 "
            "and set a HF_TOKEN environment variable."
        ) from last_error

    def diarize(
        self,
        request_id: str,
        file_path: str,
        num_speakers: int | None = None,
        min_speakers: int | None = None,
        max_speakers: int | None = None,
        hf_token: str | None = None,
    ) -> DiarizationResult:
        """Run speaker diarization on an audio file.

        Args:
            request_id: IPC request ID for progress messages.
            file_path: Path to audio file.
            num_speakers: Exact number of speakers (if known).
            min_speakers: Minimum expected speakers.
            max_speakers: Maximum expected speakers.

        Returns:
            DiarizationResult with speaker segments.
        """
        write_message(
            progress_message(request_id, 0, "loading_diarization", "Loading diarization model...")
        )

        pipeline = self._ensure_pipeline(hf_token=hf_token)

        write_message(
            progress_message(request_id, 20, "diarizing", "Running speaker diarization...")
        )

        start_time = time.time()

        # Build kwargs for speaker constraints
        kwargs: dict[str, Any] = {}
        if num_speakers is not None:
            kwargs["num_speakers"] = num_speakers
        if min_speakers is not None:
            kwargs["min_speakers"] = min_speakers
        if max_speakers is not None:
            kwargs["max_speakers"] = max_speakers

        # Convert to WAV to work around pyannote v4.0.4 duration bug
        audio_path, temp_wav = _ensure_wav(file_path)

        print(
            f"[sidecar] Running diarization on {audio_path} with kwargs: {kwargs}",
            file=sys.stderr,
            flush=True,
        )

        # Run diarization
        try:
            raw_result = pipeline(audio_path, **kwargs)
        finally:
            if temp_wav:
                os.unlink(temp_wav)

        # pyannote 4.0+ returns DiarizeOutput; older versions return Annotation directly
        if hasattr(raw_result, "speaker_diarization"):
            diarization = raw_result.speaker_diarization
        else:
            diarization = raw_result

        # Convert pyannote output to our format
        result = DiarizationResult()
        seen_speakers: set[str] = set()

        for turn, _, speaker in diarization.itertracks(yield_label=True):
            result.speaker_segments.append(
                SpeakerSegment(
                    speaker=speaker,
                    start_ms=int(turn.start * 1000),
                    end_ms=int(turn.end * 1000),
                )
            )
            seen_speakers.add(speaker)

        result.speakers = sorted(seen_speakers)
        result.num_speakers = len(seen_speakers)

        elapsed = time.time() - start_time
        print(
            f"[sidecar] Diarization complete: {result.num_speakers} speakers, "
            f"{len(result.speaker_segments)} segments in {elapsed:.1f}s",
            file=sys.stderr,
            flush=True,
        )

        write_message(
            progress_message(request_id, 100, "done", "Diarization complete")
        )

        return result


def diarization_to_payload(result: DiarizationResult) -> dict[str, Any]:
    """Convert DiarizationResult to IPC payload dict."""
    return {
        "speaker_segments": [
            {
                "speaker": seg.speaker,
                "start_ms": seg.start_ms,
                "end_ms": seg.end_ms,
            }
            for seg in result.speaker_segments
        ],
        "num_speakers": result.num_speakers,
        "speakers": result.speakers,
    }