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

feat: add diarize_model arg to CLI (#1101)

Browse Source
This commit is contained in:
bog
2025-05-31 13:32:31 +02:00
committed by GitHub
parent 6fe0a8784a
commit b343241253
3 changed files with 7 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
whisperx/diarize.py
View File

@ -11,13 +11,14 @@ from whisperx.types import TranscriptionResult, AlignedTranscriptionResult
class DiarizationPipeline:
def __init__(
self,
model_name="pyannote/speaker-diarization-3.1",
model_name=None,
use_auth_token=None,
device: Optional[Union[str, torch.device]] = "cpu",
):
if isinstance(device, str):
device = torch.device(device)
self.model = Pipeline.from_pretrained(model_name, use_auth_token=use_auth_token).to(device)
model_config = model_name or "pyannote/speaker-diarization-3.1"
self.model = Pipeline.from_pretrained(model_config, use_auth_token=use_auth_token).to(device)
def __call__(
self,