Run pyannote diarization on the GPU when CUDA is available (17.5x faster) - #51
Open
Suijiku wants to merge 1 commit into
Open
Run pyannote diarization on the GPU when CUDA is available (17.5x faster)#51Suijiku wants to merge 1 commit into
Suijiku wants to merge 1 commit into
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #50.
Summary
diarizer.pyloaded the pyannote pipeline but never moved it to a GPU, so speakerdiarization 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
DiarizationWorkerand moves the pipelineto 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 whentorch.cuda.is_available(), else"cpu"._get_pipeline(hf_token, device="cpu")— moves the pipeline viapipeline.to(torch.device("cuda"))when resolved to CUDA; cache key is now
(hf_token, resolved_device)(a CPU pipeline and aCUDA pipeline are distinct objects).
DiarizationWorkergains adeviceparam and emits a "Running speaker diarization on GPU…" note.app/main_window.py— passestranscription.devicefrom config toDiarizationWorker.tests/test_diarizer.py— device-selection tests (CUDA→GPU, CUDA-unavailable→CPU, CPU→CPU,per-device cache) with torch/pyannote mocked.
Notes
which pre-Volta GPUs handle fine — verified on the GTX 1080 below.
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:End-to-end through
DiarizationWorker(device="cuda"): emits "Running speaker diarization onGPU…", 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).