perf(pipeline): optimize embedding extraction in SpeakerDiarization

[leochiodi]

Summary

Optimizes get_embeddings in SpeakerDiarization with a fast path that activates when:

• Audio is passed in-memory (dict with waveform/sample_rate keys)
• Embedding model supports forward_frames/forward_embedding split API (e.g. WeSpeaker)

Optimizations

• Pre-slice waveforms: all chunk waveforms are sliced once upfront instead of calling audio.crop() per chunk
• Skip inactive pairs: (chunk, speaker) pairs with all-zero masks are skipped entirely (~60% fewer embedding calls on typical audio)
• Split forward pass: forward_frames runs once per chunk (shared across speakers), then forward_embedding (stats pooling only) runs per active pair
• Optional mixed precision: new embedding_precision parameter (e.g. torch.float16) enables autocast on forward_frames for GPU tensor cores, with pooling remaining in fp32 to preserve clustering quality

Backward compatibility

• Falls back to the original _get_embeddings_legacy path for file-based audio or models without the split API (SpeechBrain, NeMo, xvector)
• No change to default behavior — all optimizations are opt-in or automatically detected
• Training cache logic preserved

Benchmarks

Tested on H100 GPU with 57-minute / 10-speaker audio:

Configuration	Wall time
Original (get_embeddings)	52s
Fast path (fp32)	20s
Fast path + embedding_precision=torch.float16	~10s

The forward_frames call specifically went from 12s → 3s with float16 autocast.

Usage

Works automatically when passing in-memory audio (the common case):

pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1")
pipeline.to(torch.device("cuda"))
output = pipeline({"waveform": waveform, "sample_rate": 16000})

To enable mixed precision on the embedding forward pass:

# via config.yaml: set embedding_precision to float16
# or programmatically after loading:
pipeline._embedding_precision = torch.float16