Skip to content

Run pyannote diarization on the GPU when CUDA is available (17.5x faster) - #51

Open
Suijiku wants to merge 1 commit into
ObscureAintSecure:masterfrom
Suijiku:gpu-diarization
Open

Run pyannote diarization on the GPU when CUDA is available (17.5x faster)#51
Suijiku wants to merge 1 commit into
ObscureAintSecure:masterfrom
Suijiku:gpu-diarization

Conversation

@Suijiku

@Suijiku Suijiku commented Jul 10, 2026

Copy link
Copy Markdown

Closes #50.

Summary

diarizer.py loaded the pyannote pipeline but never moved it to a GPU, so speaker
diarization ran entirely on the CPU — even when the user selected CUDA and transcription
used the GPU. On long recordings the diarization step then dominates total time.

This threads the configured compute device into DiarizationWorker and moves the pipeline
to CUDA when available, falling back to CPU otherwise. Only the neural stages (segmentation,
embedding) move to the GPU; the final PLDA clustering stays CPU-bound regardless — but the
neural stages are the bulk of the cost on long audio.

Changes

  • app/transcription/diarizer.py
    • _resolve_device(device) — returns "cuda" only when torch.cuda.is_available(), else "cpu".
    • _get_pipeline(hf_token, device="cpu") — moves the pipeline via pipeline.to(torch.device("cuda"))
      when resolved to CUDA; cache key is now (hf_token, resolved_device) (a CPU pipeline and a
      CUDA pipeline are distinct objects).
    • DiarizationWorker gains a device param and emits a "Running speaker diarization on GPU…" note.
  • app/main_window.py — passes transcription.device from config to DiarizationWorker.
  • tests/test_diarizer.py — device-selection tests (CUDA→GPU, CUDA-unavailable→CPU, CPU→CPU,
    per-device cache) with torch/pyannote mocked.

Notes

  • No pre-Volta concern. Unlike CTranslate2 (see the compute-type PR), pyannote uses fp32,
    which pre-Volta GPUs handle fine — verified on the GTX 1080 below.
  • Portable / non-regressive. CUDA-less machines (Intel/AMD) resolve to CPU exactly as before.
  • Accuracy unchanged. CPU and GPU produce the same 4-speaker result on the test recording.

Testing

Measured on a GTX 1080 (Pascal), Windows 11, CUDA 12.6 PyTorch, pyannote.audio 4.0.7,
speaker-diarization-community-1, on a real 21.5-minute recording:

Path Wall-clock Speakers
CPU 13.32 min (799 s) 4
GPU 0.76 min (46 s) 4
Speedup 17.5× identical

End-to-end through DiarizationWorker(device="cuda"): emits "Running speaker diarization on
GPU…", completes in ~52 s (incl. load), assigns SPEAKER_00–03, no errors.

  • python -m pytest tests/ -v — full suite green (includes the new device tests).

The diarization pipeline was never moved off the CPU, so speaker diarization
ran on CPU even when CUDA was selected and transcription used the GPU. Thread
the configured compute device into DiarizationWorker and move the pyannote
pipeline to CUDA when available (pipeline.to('cuda')), falling back to CPU
otherwise. The pipeline cache key now includes the resolved device. pyannote
uses fp32, so there is no pre-Volta float16 concern here.

Measured on a GTX 1080, 21.5 min recording: CPU 13.32 min -> GPU 0.76 min
(17.5x faster), identical 4-speaker result. Closes ObscureAintSecure#50.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Run pyannote diarization on the GPU when CUDA is selected

1 participant