chrissewell/whisperX
1
0
Fork 0
mirror of https://github.com/m-bain/whisperX.git synced 2025-07-01 18:17:27 -04:00
Code Issues Packages Projects Releases Wiki Activity

add device to dia pipeline @sorgfresser

Browse Source
This commit is contained in:
Max Bain
2023-05-05 11:53:51 +01:00
parent 4e2ac4e4e9
commit 07361ba1d7
2 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
whisperx/diarize.py
View File

@ -11,6 +11,7 @@ class DiarizationPipeline:
use_auth_token=None, use_auth_token=None,
device: Optional[Union[str, torch.device]] = "cpu", device: Optional[Union[str, torch.device]] = "cpu",
): ):
self.model = Pipeline.from_pretrained(model_name, use_auth_token=use_auth_token)
if isinstance(device, str): if isinstance(device, str):
device = torch.device(device) device = torch.device(device)
self.model = Pipeline.from_pretrained(model_name, use_auth_token=use_auth_token).to(device) self.model = Pipeline.from_pretrained(model_name, use_auth_token=use_auth_token).to(device)

2
whisperx/transcribe.py
View File

@ -190,7 +190,7 @@ def cli():
tmp_results = results tmp_results = results
print(">>Performing diarization...") print(">>Performing diarization...")
results = [] results = []
diarize_model = DiarizationPipeline(use_auth_token=hf_token) diarize_model = DiarizationPipeline(use_auth_token=hf_token, device=device)
for result, input_audio_path in tmp_results: for result, input_audio_path in tmp_results:
diarize_segments = diarize_model(input_audio_path, min_speakers=min_speakers, max_speakers=max_speakers) diarize_segments = diarize_model(input_audio_path, min_speakers=min_speakers, max_speakers=max_speakers)
results_segments, word_segments = assign_word_speakers(diarize_segments, result["segments"]) results_segments, word_segments = assign_word_speakers(diarize_segments, result["segments"])