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

allow custom model_dir for torchaudio models

Browse Source
This commit is contained in:
Max Bain
2023-04-14 21:40:36 +01:00
parent cf252a8592
commit da458863d7
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
whisperx/alignment.py
View File

@ -40,7 +40,7 @@ DEFAULT_ALIGN_MODELS_HF = {
} }
def load_align_model(language_code, device, model_name=None): def load_align_model(language_code, device, model_name=None, model_dir=None):
if model_name is None: if model_name is None:
# use default model # use default model
if language_code in DEFAULT_ALIGN_MODELS_TORCH: if language_code in DEFAULT_ALIGN_MODELS_TORCH:
@ -55,7 +55,7 @@ def load_align_model(language_code, device, model_name=None):
if model_name in torchaudio.pipelines.__all__: if model_name in torchaudio.pipelines.__all__:
pipeline_type = "torchaudio" pipeline_type = "torchaudio"
bundle = torchaudio.pipelines.__dict__[model_name] bundle = torchaudio.pipelines.__dict__[model_name]
align_model = bundle.get_model().to(device) align_model = bundle.get_model(dl_kwargs={"model_dir": model_dir}).to(device)
labels = bundle.get_labels() labels = bundle.get_labels()
align_dictionary = {c.lower(): i for i, c in enumerate(labels)} align_dictionary = {c.lower(): i for i, c in enumerate(labels)}
else: else: