Speaker Diarization with Pyannote

Within the Transcription & Speaker Diarization pipeline, speaker diarization with Pyannote is the attribution boundary that converts a continuous acoustic waveform into temporally aligned, speaker-clustered metadata — the “who spoke when” layer that the rest of the system contracts against. It protects a single, unforgiving SLA: every transcript segment must carry a stable speaker label whose boundaries align to the same timeline the acoustic model emits, because once attribution drifts, multi-host podcasts, compliance-driven captioning, and interview archives all corrupt downstream in ways that are expensive to debug. This page covers running pyannote.audio as a deterministic, memory-bounded stage rather than an isolated inference call.

Prerequisites & Environment

The diarization stage assumes Python 3.10+ (for precise type hints on the segment models and structural pattern matching on event types) and a CUDA-capable GPU with at least 4 GB of free VRAM for the default 15-second window configuration. The pyannote/speaker-diarization-3.1 pipeline is gated on the Hugging Face Hub, so the worker container needs a read token exported as HF_TOKEN and must accept the model’s user conditions before the first pull. Pin every weight: the segmentation model, the embedding model, and the pipeline definition all version independently, and an unpinned cache silently degrades accuracy on a routine re-pull.

# Python 3.10+
# pyannote.audio==3.1.1     # diarization pipeline (segmentation + embedding + clustering)
# torch==2.2.1              # inference backend; match CUDA build to the host driver
# torchaudio==2.2.1         # resampling + waveform I/O
# soundfile==0.12.1         # deterministic WAV read/write
# pydantic==2.6.4           # output-contract validation
# structlog==24.1.0         # JSON structured logging
# prometheus-client==0.20.0 # VRAM + segment-count metrics

Workers running diarization must be GPU-affinity isolated from the CPU-bound preprocessing pool so that resampling and channel-mixing never compete with embedding extraction for VRAM. Everything dynamic lives in environment variables; static routing keys and model revisions stay version-controlled in the worker config, not the shell.

This stage consumes the 16 kHz mono, 16-bit PCM audio guaranteed by the audio codec normalization workflows in the Media Ingestion & Format Architecture boundary, so the diarizer can assume a clean, fetchable waveform and spend its cycles on attribution rather than format repair.

Architecture & Queue Topology

Diarization runs as a parallel consumer rather than a serial blocker. The async transcription queue leases a task that references a normalized audio object, and the diarization worker fans out alongside the acoustic model — both read the same waveform from object storage, neither waits on the other, and their outputs converge later at the alignment boundary. Heavy audio never traverses the broker; the task body carries only a task_id, the immutable source_uri, and a presigned fetch URL. The worker pulls the waveform, runs chunked inference, serializes the turns, and publishes a completion event that the alignment and editorial consumers subscribe to.

Internally the Pipeline class collapses segmentation, embedding extraction, agglomerative clustering, and overlap resolution behind one callable. Treating it as a black box is what causes production incidents: naive full-file inference on 60-minute media routinely triggers CUDA out of memory on shared GPUs. The worker therefore drives the pipeline over overlapping windows it controls, so VRAM ceilings stay predictable regardless of source duration.

Step-by-Step Implementation

Each step ends with a verification command you can run before wiring the next one.

1. Pin the model and authenticate the cache. Load the pipeline by exact revision and move it to the GPU once, at worker startup, not per task:

# Python 3.10+
import os
import torch
from pyannote.audio import Pipeline

PIPELINE = Pipeline.from_pretrained(
    os.environ["PYANNOTE_REVISION"],          # pinned: pyannote/speaker-diarization-3.1
    use_auth_token=os.environ["HF_TOKEN"],
)
PIPELINE.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))

Verify the weights resolved to the pinned revision and landed on the GPU:

python -c "from pyannote.audio import Pipeline; import os; \
p = Pipeline.from_pretrained(os.environ['PYANNOTE_REVISION'], use_auth_token=os.environ['HF_TOKEN']); \
print('loaded', p.parameters() is not None)"

2. Validate the audio input contract. Reject anything that is not 16 kHz, mono, 16-bit PCM at the ingress instead of letting a sample-rate mismatch corrupt every downstream timestamp:

# Python 3.10+
import soundfile as sf

def assert_contract(path: str) -> None:
    info = sf.info(path)
    if info.samplerate != 16000:
        raise ValueError(f"expected 16000 Hz, got {info.samplerate}")
    if info.channels != 1:
        raise ValueError(f"expected mono, got {info.channels} channels")
    if info.subtype != "PCM_16":
        raise ValueError(f"expected PCM_16, got {info.subtype}")

Verify against a sample before inference:

ffprobe -v error -show_entries stream=sample_rate,channels,sample_fmt \
  -of default=noprint_wrappers=1 sample.wav

3. Run memory-aware chunked inference. Drive the pipeline over overlapping windows and poll VRAM before each chunk, dropping batch size and falling back to CPU embedding extraction when free memory crosses the floor:

# Python 3.10+
import gc
import torch

WINDOW = float(os.environ["DIAR_WINDOW_SECONDS"])   # 15.0
STRIDE = float(os.environ["DIAR_STRIDE_SECONDS"])    # 3.0
VRAM_FLOOR = int(os.environ["DIAR_VRAM_FLOOR_MB"]) * 1024 * 1024

def free_vram() -> int:
    free, _ = torch.cuda.mem_get_info()
    return free

def diarize(path: str, min_speakers: int, max_speakers: int):
    if torch.cuda.is_available() and free_vram() < VRAM_FLOOR:
        PIPELINE.to(torch.device("cpu"))            # graceful degradation, not OOM
    annotation = PIPELINE(
        path,
        min_speakers=min_speakers,
        max_speakers=max_speakers,
    )
    torch.cuda.empty_cache()                         # release fragmented blocks
    gc.collect()
    return annotation

The min_speakers and max_speakers bounds are load-bearing: passing a known speaker count prevents the speaker-merging collapse described under failure modes. Verify a clip diarizes without an OOM:

python -c "from worker import diarize; \
ann = diarize('sample.wav', 2, 4); print(len(list(ann.itertracks())), 'turns')"

4. Serialize turns to RTTM and JSON. Round every boundary to 10 ms precision and validate the JSON payload against a fixed schema before it leaves the worker:

# Python 3.10+
from pydantic import BaseModel, Field

class Turn(BaseModel):
    start_time: float = Field(ge=0)
    end_time: float = Field(ge=0)
    speaker_id: str
    confidence_score: float = Field(ge=0, le=1)

def to_payload(annotation, scores) -> list[dict]:
    turns: list[Turn] = []
    for segment, _, label in annotation.itertracks(yield_label=True):
        turns.append(Turn(
            start_time=round(segment.start, 2),
            end_time=round(segment.end, 2),
            speaker_id=label,
            confidence_score=round(scores.get(label, 1.0), 4),
        ))
    return [t.model_dump() for t in turns]

Pyannote also writes canonical RTTM directly, which is the interchange format graders and review tooling expect:

python -c "from worker import diarize; \
ann = diarize('sample.wav', 2, 4); \
open('sample.rttm','w').write(ann.to_rttm())" && head -n 3 sample.rttm

5. Publish the completion event downstream. Emit a single event carrying the validated turns, the input hash, and chunk metadata so the alignment consumer can merge attribution with lexical tokens:

# Python 3.10+
import hashlib, json

def completion_event(path: str, payload: list[dict]) -> str:
    digest = hashlib.sha256(open(path, "rb").read()).hexdigest()
    return json.dumps({
        "task_id": os.environ.get("TASK_ID"),
        "source_sha256": digest,
        "turns": payload,
        "model_revision": os.environ["PYANNOTE_REVISION"],
    })

Verify the event is well-formed JSON and references the expected revision:

python -c "from worker import completion_event; \
import json; e = json.loads(completion_event('sample.wav', [])); \
print(e['model_revision'])"

Data Contracts

The diarizer consumes a normalized waveform and produces a speaker-turn payload. Both sides are validated; a missing or malformed field triggers immediate task rejection rather than silent degradation, so downstream consumers never ingest a partial attribution map.

RTTM and JSON are emitted together: RTTM for review tooling and benchmark scoring, JSON for the queue-routed Whisper Large V3 integration consumer that fuses speaker clusters with lexical tokens.

Resilience Patterns

Diarization is expensive and non-idempotent if rerun carelessly, so failure handling is explicit rather than incidental. Tasks are keyed on a stable task_id and outputs are written to a content-addressed path derived from the source_sha256; a redelivered message under at-least-once delivery that finds a completed result re-acknowledges instead of re-running inference. Transient faults — a presigned URL that expired mid-fetch, a recoverable CUDA context error — are retried with exponential backoff and jitter up to MAX_TASK_ATTEMPTS, after which the task routes to a dead-letter queue rather than blocking the lane.

This stage leans on the same retry logic and dead-letter queues the rest of the batch layer uses: a CUDA out of memory failure is retried once at a smaller window before being dead-lettered, while a contract violation (wrong sample rate, corrupt waveform) is dead-lettered immediately because retrying identical bad input only burns GPU time. Cross-host placement of these GPU and CPU pools is governed by Celery task routing, which pins diarization tasks to GPU-affinity queues so they never land on a preprocessing worker.

Observability & Debugging

Instrument every pipeline boundary so attribution regressions surface before a consumer complains. The metrics that matter are VRAM headroom at chunk entry, segment count per minute of audio, mean cluster confidence, and the rate of low-confidence segments routed to review. Structured logs should capture the input hash, chunk dimensions, the resolved model_revision, and the schema-validation result for every task.

Cluster collapse and timestamp drift are the two failures worth alerting on directly. Collapse is caught by validating the produced speaker count against the known host count and firing when they disagree; drift is caught by asserting that the final end_time never exceeds the waveform duration read from ffprobe. For VRAM fragmentation on long-running workers, the official PyTorch CUDA memory management notes cover torch.cuda.empty_cache() behavior and the PYTORCH_CUDA_ALLOC_CONF knobs that mitigate it; the canonical pipeline reference lives in the pyannote.audio documentation.

Performance Tuning

GPU workers run diarization at a concurrency of exactly one — a second concurrent task on the same device races for VRAM and triggers the OOM-kill that takes the whole worker down mid-file. Scale horizontally by adding workers, not by raising per-worker concurrency. The 15-second window with a 3-second stride is the default balance between temporal resolution and VRAM consumption; shorter windows raise boundary precision at the cost of more embedding passes, while longer windows reduce overhead but raise the OOM risk on consumer-grade cards.

For degraded source material — compressed VoIP or heavily reverberant room audio — raise DIAR_CONF_THRESHOLD to suppress spurious clusters and route the low-confidence remainder to manual review rather than letting noise pollute the attribution map. Cost-optimized routing then directs high-confidence diarized chunks to the cheaper transcription endpoint while reserving the premium acoustic model for overlapping speech and low-SNR segments. After inference, the worker’s refined boundaries flow into timestamp alignment, where they are fused with word-level tokens to sub-100 ms precision before the editorial layer in the Pipeline Automation & Batch Processing graph consumes them.

Frequently Asked Questions

Related

• Up to the parent overview: Transcription & Speaker Diarization
• The queue that leases diarization tasks: async transcription queue management
• The parallel acoustic-to-text engine: Whisper Large V3 integration
• Where diarization boundaries get refined: timestamp alignment & correction
• The normalized input this stage assumes: audio codec normalization workflows
• Where failed diarization tasks are quarantined: retry logic and dead-letter queues