src/diarization.py

from pathlib import Path
from typing import Tuple

import torch
from pyannote.audio import Pipeline
from pydub import AudioSegment


def get_pipeline(filename: str | Path, hf_api_key: str, pyannote_model: str, tmp_dir: Path) -> Tuple[AudioSegment, Pipeline]:
    pipeline = Pipeline.from_pretrained(
        pyannote_model,
        token=hf_api_key,
    )
    pipeline.to(torch.device("cuda"))

    audio_segment = AudioSegment.from_mp3(filename)
    wav_audio = tmp_dir.joinpath(Path(filename).name).with_suffix(".wav")

    with open(wav_audio, "wb"):
        audio_segment.export(wav_audio, format="wav")

    return (audio_segment, pipeline(wav_audio))