add device to dia pipeline @sorgfresser

m-bain · May 5, 2023 · 07361ba · 07361ba · sorgfresser · May 5, 2023
1 parent 4e2ac4e
commit 07361ba
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/whisperx/diarize.py b/whisperx/diarize.py
@@ -11,6 +11,7 @@ def __init__(
         use_auth_token=None,
         device: Optional[Union[str, torch.device]] = "cpu",
     ):
+        self.model = Pipeline.from_pretrained(model_name, use_auth_token=use_auth_token)
         if isinstance(device, str):
             device = torch.device(device)
         self.model = Pipeline.from_pretrained(model_name, use_auth_token=use_auth_token).to(device)

diff --git a/whisperx/transcribe.py b/whisperx/transcribe.py
@@ -190,7 +190,7 @@ def cli():
         tmp_results = results
         print(">>Performing diarization...")
         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:
             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"])