Merge pull request #1469 from li-lizhe/fix-optional-torchaudio

Make torchaudio optional with Hugging Face fallback for align models

- Tests look good, happy with merge

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"

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 = {"<pad>": 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

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"

whisperx/alignment.py

@@ -8,7 +8,10 @@ from typing import Iterable, Optional, Union, List
 import numpy as np
 import pandas as pd
 import torch
+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)