diff --git a/pyproject.toml b/pyproject.toml index fdde305..bc7af14 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,7 +18,6 @@ dependencies = [ "pyannote-audio>=4.0.0", "huggingface-hub>=0.28.1", "torch~=2.8.0", - "torchaudio~=2.8.0", "torchvision~=0.23.0", "torchcodec>=0.6.0,<0.8.0; (sys_platform == 'linux' and platform_machine == 'x86_64') or sys_platform == 'darwin' or sys_platform == 'win32'", "transformers>=4.48.0", @@ -29,6 +28,10 @@ dependencies = [ [project.optional-dependencies] dev = ["pytest"] fast = ["numba>=0.68.0"] +# torchaudio is only used for loading align models via torchaudio.pipelines +# (an optional path). Platforms without torchaudio wheels (e.g. Ascend / +# Linux aarch64) skip installing it and use the HF fallback instead. +align-torch = ["torchaudio~=2.8.0"] [project.scripts] whisperx = "whisperx.__main__:cli" diff --git a/tests/test_align_model_fallback.py b/tests/test_align_model_fallback.py new file mode 100644 index 0000000..bd21438 --- /dev/null +++ b/tests/test_align_model_fallback.py @@ -0,0 +1,48 @@ +"""Tests for torchaudio-optional alignment model loading.""" + +from unittest import mock + +import whisperx.alignment as A + + +def test_default_torch_models_have_hf_fallback(): + """Every default torchaudio.pipelines align model should have an HF + equivalent, so platforms without torchaudio wheels (e.g. Ascend / + Linux aarch64) can still run forced alignment via Hugging Face.""" + assert set(A.DEFAULT_ALIGN_MODELS_TORCH.values()) == set( + A.TORCHAUDIO_PIPELINE_TO_HF.keys() + ) + + +def test_hf_fallback_names_are_valid_checkpoints(): + """HF fallback names should be official facebook wav2vec2 checkpoints.""" + for name in A.TORCHAUDIO_PIPELINE_TO_HF.values(): + assert name.startswith("facebook/"), name + + +def test_load_align_model_falls_back_to_hf_when_torchaudio_missing(): + """With torchaudio absent, a torchaudio.pipelines default model name must + be transparently remapped to its HF checkpoint and loaded via + Wav2Vec2 (the HuggingFace path), not crash.""" + fake_vocab = {"": 0, "a": 1, "b": 2} + fake_processor = mock.MagicMock() + fake_processor.tokenizer.get_vocab.return_value = fake_vocab + fake_model = mock.MagicMock() + + with mock.patch.object(A, "torchaudio", None), mock.patch.object( + A, "Wav2Vec2Processor" + ) as proc_cls, mock.patch.object(A, "Wav2Vec2ForCTC") as model_cls: + proc_cls.from_pretrained.return_value = fake_processor + model_cls.from_pretrained.return_value = fake_model + + model, metadata = A.load_align_model( + language_code="en", device="cpu" + ) + + assert model is fake_model.to.return_value # returned after .to(device) + assert metadata["type"] == "huggingface" + assert metadata["dictionary"] == {c.lower(): i for c, i in fake_vocab.items()} + # the torchaudio bundle name was remapped to the HF checkpoint + remapped_name = proc_cls.from_pretrained.call_args.args[0] + assert remapped_name == A.TORCHAUDIO_PIPELINE_TO_HF["WAV2VEC2_ASR_BASE_960H"] + assert model_cls.from_pretrained.call_args.args[0] == remapped_name diff --git a/uv.lock b/uv.lock index cf95be4..ab317a2 100644 --- a/uv.lock +++ b/uv.lock @@ -3206,9 +3206,6 @@ dependencies = [ { name = "torch", version = "2.8.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "sys_platform == 'darwin'" }, { name = "torch", version = "2.8.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "platform_machine != 'AMD64' and platform_machine != 'x86_64' and sys_platform != 'darwin'" }, { name = "torch", version = "2.8.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "(platform_machine == 'AMD64' and sys_platform != 'darwin') or (platform_machine == 'x86_64' and sys_platform != 'darwin')" }, - { name = "torchaudio", version = "2.8.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or sys_platform == 'darwin'" }, - { name = "torchaudio", version = "2.8.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(platform_machine != 'AMD64' and platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux') or (platform_machine != 'AMD64' and platform_machine != 'x86_64' and sys_platform != 'darwin' and sys_platform != 'linux') or (platform_machine == 'aarch64' and platform_python_implementation != 'CPython' and sys_platform == 'linux')" }, - { name = "torchaudio", version = "2.8.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "(platform_machine == 'AMD64' and sys_platform != 'darwin') or (platform_machine == 'x86_64' and sys_platform != 'darwin')" }, { name = "torchcodec", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin' or sys_platform == 'win32'" }, { name = "torchvision" }, { name = "transformers" }, @@ -3216,6 +3213,11 @@ dependencies = [ ] [package.optional-dependencies] +align-torch = [ + { name = "torchaudio", version = "2.8.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or sys_platform == 'darwin'" }, + { name = "torchaudio", version = "2.8.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(platform_machine != 'AMD64' and platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux') or (platform_machine != 'AMD64' and platform_machine != 'x86_64' and sys_platform != 'darwin' and sys_platform != 'linux') or (platform_machine == 'aarch64' and platform_python_implementation != 'CPython' and sys_platform == 'linux')" }, + { name = "torchaudio", version = "2.8.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "(platform_machine == 'AMD64' and sys_platform != 'darwin') or (platform_machine == 'x86_64' and sys_platform != 'darwin')" }, +] dev = [ { name = "pytest" }, ] @@ -3238,15 +3240,15 @@ requires-dist = [ { name = "torch", marker = "sys_platform == 'darwin'", specifier = "~=2.8.0", index = "https://download.pytorch.org/whl/cpu" }, { name = "torch", marker = "platform_machine != 'AMD64' and platform_machine != 'x86_64' and sys_platform != 'darwin'", specifier = "~=2.8.0", index = "https://download.pytorch.org/whl/cpu" }, { name = "torch", marker = "(platform_machine == 'AMD64' and sys_platform != 'darwin') or (platform_machine == 'x86_64' and sys_platform != 'darwin')", specifier = "~=2.8.0", index = "https://download.pytorch.org/whl/cu128" }, - { name = "torchaudio", marker = "sys_platform == 'darwin'", specifier = "~=2.8.0", index = "https://download.pytorch.org/whl/cpu" }, - { name = "torchaudio", marker = "platform_machine != 'AMD64' and platform_machine != 'x86_64' and sys_platform != 'darwin'", specifier = "~=2.8.0", index = "https://download.pytorch.org/whl/cpu" }, - { name = "torchaudio", marker = "(platform_machine == 'AMD64' and sys_platform != 'darwin') or (platform_machine == 'x86_64' and sys_platform != 'darwin')", specifier = "~=2.8.0", index = "https://download.pytorch.org/whl/cu128" }, + { name = "torchaudio", marker = "sys_platform == 'darwin' and extra == 'align-torch'", specifier = "~=2.8.0", index = "https://download.pytorch.org/whl/cpu" }, + { name = "torchaudio", marker = "platform_machine != 'AMD64' and platform_machine != 'x86_64' and sys_platform != 'darwin' and extra == 'align-torch'", specifier = "~=2.8.0", index = "https://download.pytorch.org/whl/cpu" }, + { name = "torchaudio", marker = "((platform_machine == 'AMD64' and sys_platform != 'darwin') or (platform_machine == 'x86_64' and sys_platform != 'darwin')) and extra == 'align-torch'", specifier = "~=2.8.0", index = "https://download.pytorch.org/whl/cu128" }, { name = "torchcodec", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin' or sys_platform == 'win32'", specifier = ">=0.6.0,<0.8.0" }, { name = "torchvision", specifier = "~=0.23.0" }, { name = "transformers", specifier = ">=4.48.0" }, { name = "triton", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'", specifier = ">=3.3.0", index = "https://download.pytorch.org/whl/cu128" }, ] -provides-extras = ["dev", "fast"] +provides-extras = ["align-torch", "dev", "fast"] [[package]] name = "yarl" diff --git a/whisperx/alignment.py b/whisperx/alignment.py index 2987c36..378a24f 100644 --- a/whisperx/alignment.py +++ b/whisperx/alignment.py @@ -8,7 +8,10 @@ from typing import Iterable, Optional, Union, List import numpy as np import pandas as pd import torch -import torchaudio +try: + import torchaudio +except Exception: # broken installs (missing native libs, ABI mismatch) should + torchaudio = None # degrade to the HuggingFace path, not crash the import from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor try: @@ -54,6 +57,17 @@ DEFAULT_ALIGN_MODELS_TORCH = { "it": "VOXPOPULI_ASR_BASE_10K_IT", } +# Hugging Face equivalents of the torchaudio.pipelines checkpoints. +# torchaudio has no Linux aarch64 wheels (Ascend/Kunpeng ARM etc. cannot install +# it), so fall back to these HF models to keep forced alignment available. +TORCHAUDIO_PIPELINE_TO_HF = { + "WAV2VEC2_ASR_BASE_960H": "facebook/wav2vec2-base-960h", + "VOXPOPULI_ASR_BASE_10K_FR": "facebook/wav2vec2-base-10k-voxpopuli-fr", + "VOXPOPULI_ASR_BASE_10K_DE": "facebook/wav2vec2-base-10k-voxpopuli-de", + "VOXPOPULI_ASR_BASE_10K_ES": "facebook/wav2vec2-base-10k-voxpopuli-es", + "VOXPOPULI_ASR_BASE_10K_IT": "facebook/wav2vec2-base-10k-voxpopuli-it", +} + DEFAULT_ALIGN_MODELS_HF = { "ja": "jonatasgrosman/wav2vec2-large-xlsr-53-japanese", "zh": "jonatasgrosman/wav2vec2-large-xlsr-53-chinese-zh-cn", @@ -107,13 +121,17 @@ def load_align_model(language_code: str, device: str, model_name: Optional[str] f"then pass the model name via --align_model [MODEL_NAME]") raise ValueError(f"No default align-model for language: {language_code}") - if model_name in torchaudio.pipelines.__all__: + if torchaudio is not None and model_name in torchaudio.pipelines.__all__: pipeline_type = "torchaudio" bundle = torchaudio.pipelines.__dict__[model_name] align_model = bundle.get_model(dl_kwargs={"model_dir": model_dir}).to(device) labels = bundle.get_labels() align_dictionary = {c.lower(): i for i, c in enumerate(labels)} else: + # When torchaudio is unavailable, map torchaudio.pipelines model names + # to their equivalent HF checkpoints. + if torchaudio is None and model_name in TORCHAUDIO_PIPELINE_TO_HF: + model_name = TORCHAUDIO_PIPELINE_TO_HF[model_name] try: processor = Wav2Vec2Processor.from_pretrained(model_name, cache_dir=model_dir, local_files_only=model_cache_only) align_model = Wav2Vec2ForCTC.from_pretrained(model_name, cache_dir=model_dir, local_files_only=model_cache_only)