Merge 399010fd12 into d700b56c9c

This reverts commit 6fe0a8784a.

Revert the commit now that the issue is fixed.

Signed-off-by: CHEN, CHUN <jim60105@gmail.com>

.github/workflows/build-and-release.yml

@@ -17,6 +17,9 @@ jobs:
           version: "0.5.14"
           python-version: "3.9"
 
+      - name: Check if lockfile is up to date
+        run: uv lock --check
+
       - name: Build package
         run: uv build
 

.github/workflows/python-compatibility.yml

@@ -23,6 +23,9 @@ jobs:
           version: "0.5.14"
           python-version: ${{ matrix.python-version }}
 
+      - name: Check if lockfile is up to date
+        run: uv lock --check
+
       - name: Install the project
         run: uv sync --all-extras
 

README.md

@@ -170,7 +170,7 @@ result = model.transcribe(audio, batch_size=batch_size)
 print(result["segments"]) # before alignment
 
 # delete model if low on GPU resources
-# import gc; gc.collect(); torch.cuda.empty_cache(); del model
+# import gc; import torch; gc.collect(); torch.cuda.empty_cache(); del model
 
 # 2. Align whisper output
 model_a, metadata = whisperx.load_align_model(language_code=result["language"], device=device)
@@ -179,7 +179,7 @@ result = whisperx.align(result["segments"], model_a, metadata, audio, device, re
 print(result["segments"]) # after alignment
 
 # delete model if low on GPU resources
-# import gc; gc.collect(); torch.cuda.empty_cache(); del model_a
+# import gc; import torch; gc.collect(); torch.cuda.empty_cache(); del model_a
 
 # 3. Assign speaker labels
 diarize_model = whisperx.diarize.DiarizationPipeline(use_auth_token=YOUR_HF_TOKEN, device=device)

pyproject.toml

@@ -13,11 +13,11 @@ dependencies = [
     "faster-whisper>=1.1.1",
     "nltk>=3.9.1",
     "numpy>=2.0.2",
-    "onnxruntime>=1.19",
+    "onnxruntime>=1.19,<1.20.0",
     "pandas>=2.2.3",
     "pyannote-audio>=3.3.2",
-    "torch>=2.5.1",
-    "torchaudio>=2.5.1",
+    "torch<2.4.0",
+    "torchaudio",
     "transformers>=4.48.0",
 ]
 

whisperx/__main__.py

@@ -43,6 +43,7 @@ def cli():
     parser.add_argument("--diarize", action="store_true", help="Apply diarization to assign speaker labels to each segment/word")
     parser.add_argument("--min_speakers", default=None, type=int, help="Minimum number of speakers to in audio file")
     parser.add_argument("--max_speakers", default=None, type=int, help="Maximum number of speakers to in audio file")
+    parser.add_argument("--diarize_model", default="pyannote/speaker-diarization-3.1", type=str, help="Name of the speaker diarization model to use")
 
     parser.add_argument("--temperature", type=float, default=0, help="temperature to use for sampling")
     parser.add_argument("--best_of", type=optional_int, default=5, help="number of candidates when sampling with non-zero temperature")

whisperx/diarize.py

@@ -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,

whisperx/transcribe.py

@@ -57,6 +57,7 @@ def transcribe_task(args: dict, parser: argparse.ArgumentParser):
     diarize: bool = args.pop("diarize")
     min_speakers: int = args.pop("min_speakers")
     max_speakers: int = args.pop("max_speakers")
+    diarize_model_name: str = args.pop("diarize_model")
     print_progress: bool = args.pop("print_progress")
 
     if args["language"] is not None:
@@ -204,8 +205,9 @@ def transcribe_task(args: dict, parser: argparse.ArgumentParser):
             )
         tmp_results = results
         print(">>Performing diarization...")
+        print(">>Using model:", diarize_model_name)
         results = []
-        diarize_model = DiarizationPipeline(use_auth_token=hf_token, device=device)
+        diarize_model = DiarizationPipeline(model_name=diarize_model_name, 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