Diarization now supports version selection, with the default set to speaker-diarization-3.1.

app.py

@@ -19,7 +19,6 @@ from src.diarization.diarization import Diarization
 from src.diarization.diarizationContainer import DiarizationContainer
 from src.hooks.progressListener import ProgressListener
 from src.hooks.subTaskProgressListener import SubTaskProgressListener
-from src.hooks.whisperProgressHook import create_progress_listener_handle
 from src.modelCache import ModelCache
 from src.prompts.jsonPromptStrategy import JsonPromptStrategy
 from src.prompts.prependPromptStrategy import PrependPromptStrategy
@@ -32,7 +31,7 @@ import ffmpeg
 # UI
 import gradio as gr
 
-from src.download import ExceededMaximumDuration, download_url
+from src.download import ExceededMaximumDuration
 from src.utils import optional_int, slugify, str2bool, write_srt, write_srt_original, write_vtt
 from src.vad import AbstractTranscription, NonSpeechStrategy, PeriodicTranscriptionConfig, TranscriptionConfig, VadPeriodicTranscription, VadSileroTranscription
 from src.whisper.abstractWhisperContainer import AbstractWhisperContainer
@@ -100,11 +99,15 @@ class WhisperTranscriber:
             self.vad_cpu_cores = min(os.cpu_count(), MAX_AUTO_CPU_CORES)
             print("[Auto parallel] Using GPU devices " + str(self.parallel_device_list) + " and " + str(self.vad_cpu_cores) + " CPU cores for VAD/transcription.")
 
-    def set_diarization(self, auth_token: str, enable_daemon_process: bool = True, **kwargs):
+    def set_diarization(self, auth_token: str, enable_daemon_process: bool = True, diarization_version: str = None, **kwargs):
+        if diarization_version == None:
+            diarization_version = self.app_config.diarization_version
         if self.diarization is None:
             self.diarization = DiarizationContainer(auth_token=auth_token, enable_daemon_process=enable_daemon_process, 
                                                     auto_cleanup_timeout_seconds=self.app_config.diarization_process_timeout, 
-                                                    cache=self.model_cache)
+                                                    cache=self.model_cache, diarization_version=diarization_version)
+        else:
+            self.diarization.diarization_version=diarization_version
         # Set parameters
         self.diarization_kwargs = kwargs
 
@@ -257,6 +260,7 @@ class WhisperTranscriber:
             diarization_speakers:     int  = decodeOptions.pop("diarization_speakers", 2)
             diarization_min_speakers: int  = decodeOptions.pop("diarization_min_speakers", 1)
             diarization_max_speakers: int  = decodeOptions.pop("diarization_max_speakers", 8)
+            diarization_version:      str  = decodeOptions.pop("diarization_version", "speaker-diarization-3.1")
             highlight_words:          bool = decodeOptions.pop("highlight_words", False)
             
             temperature: float = decodeOptions.pop("temperature", None)
@@ -290,9 +294,9 @@ class WhisperTranscriber:
 
             if diarization:
                 if diarization_speakers is not None and diarization_speakers < 1:
-                    self.set_diarization(auth_token=self.app_config.auth_token, min_speakers=diarization_min_speakers, max_speakers=diarization_max_speakers)
+                    self.set_diarization(auth_token=self.app_config.auth_token, min_speakers=diarization_min_speakers, max_speakers=diarization_max_speakers, diarization_version=diarization_version)
                 else:
-                    self.set_diarization(auth_token=self.app_config.auth_token, num_speakers=diarization_speakers, min_speakers=diarization_min_speakers, max_speakers=diarization_max_speakers)
+                    self.set_diarization(auth_token=self.app_config.auth_token, num_speakers=diarization_speakers, min_speakers=diarization_min_speakers, max_speakers=diarization_max_speakers, diarization_version=diarization_version)
             else:
                 self.unset_diarization()
                 
@@ -1137,7 +1141,8 @@ def create_ui(app_config: ApplicationConfig):
         gr.Checkbox(label="Diarization", value=app_config.diarization, interactive=has_diarization_libs, elem_id="diarization", info="Whether to perform speaker diarization"),
         gr.Number(label="Diarization - Speakers", precision=0, value=app_config.diarization_speakers, interactive=has_diarization_libs, elem_id="diarization_speakers", info="The number of speakers to detect"),
         gr.Number(label="Diarization - Min Speakers", precision=0, value=app_config.diarization_min_speakers, interactive=has_diarization_libs, elem_id="diarization_min_speakers", info="The minimum number of speakers to detect"),
-        gr.Number(label="Diarization - Max Speakers", precision=0, value=app_config.diarization_max_speakers, interactive=has_diarization_libs, elem_id="diarization_max_speakers", info="The maximum number of speakers to detect")
+        gr.Number(label="Diarization - Max Speakers", precision=0, value=app_config.diarization_max_speakers, interactive=has_diarization_libs, elem_id="diarization_max_speakers", info="The maximum number of speakers to detect"),
+        gr.Dropdown(label="Diarization Version", choices=["speaker-diarization-3.1", "speaker-diarization-3.0", "[email protected]"], value=app_config.diarization_version, elem_id="diarization_version", info="pyannote.audio speaker diarization pipeline v3.1 is expected to be much better (and faster) than v2.x. [Benchmark](https://github.com/pyannote/pyannote-audio?tab=readme-ov-file#benchmark)"),
     }
     
     common_output = lambda : [
@@ -1439,6 +1444,7 @@ if __name__ == '__main__':
     parser.add_argument("--diarization_max_speakers", type=int, default=default_app_config.diarization_max_speakers, help="Maximum number of speakers")
     parser.add_argument("--diarization_process_timeout", type=int, default=default_app_config.diarization_process_timeout, \
                         help="Number of seconds before inactivate diarization processes are terminated. Use 0 to close processes immediately, or None for no timeout.")
+    parser.add_argument('--diarization_version', type=str, default=default_app_config.diarization_version, help='Specify the diarization version, defaulting to speaker-diarization-3.1')
 
     args = parser.parse_args().__dict__
 

requirements-fasterWhisper.txt

@@ -20,7 +20,7 @@ sentencepiece
 # Needed by diarization
 intervaltree
 srt
-https://github.com/pyannote/pyannote-audio/archive/refs/heads/develop.zip
+pyannote.audio
 
 # Needed by ALMA-GPTQ
 accelerate

requirements-whisper.txt

@@ -20,7 +20,7 @@ sentencepiece
 # Needed by diarization
 intervaltree
 srt
-https://github.com/pyannote/pyannote-audio/archive/refs/heads/develop.zip
+pyannote.audio
 
 # Needed by ALMA-GPTQ
 accelerate

requirements.txt

@@ -20,7 +20,7 @@ sentencepiece
 # Needed by diarization
 intervaltree
 srt
-https://github.com/pyannote/pyannote-audio/archive/refs/heads/develop.zip
+pyannote.audio
 
 # Needed by ALMA-GPTQ
 accelerate

src/config.py

@@ -78,6 +78,7 @@ class ApplicationConfig:
                  auth_token: str = None, diarization: bool = False, diarization_speakers: int = 2,
                  diarization_min_speakers: int = 1, diarization_max_speakers: int = 5,
                  diarization_process_timeout: int = 60,
+                 diarization_version: str = "speaker-diarization-3.1",
                  # Translation
                  translation_batch_size: int = 2,
                  translation_no_repeat_ngram_size: int = 4,
@@ -148,6 +149,8 @@ class ApplicationConfig:
         self.diarization_min_speakers = diarization_min_speakers
         self.diarization_max_speakers = diarization_max_speakers
         self.diarization_process_timeout = diarization_process_timeout
+        self.diarization_version = diarization_version
+
         # Translation
         self.translation_batch_size = translation_batch_size
         self.translation_no_repeat_ngram_size = translation_no_repeat_ngram_size

src/diarization/diarization.py

@@ -26,15 +26,16 @@ class DiarizationEntry:
         }
 
 class Diarization:
-    def __init__(self, auth_token=None):
+    def __init__(self, auth_token=None, diarization_version=None):
         if auth_token is None:
             auth_token = os.environ.get("HF_ACCESS_TOKEN")
             if auth_token is None:
                 raise ValueError("No HuggingFace API Token provided - please use the --auth_token argument or set the HF_ACCESS_TOKEN environment variable")
         
-        self.auth_token = auth_token
-        self.initialized = False
-        self.pipeline = None
+        self.auth_token          = auth_token
+        self.initialized         = False
+        self.pipeline            = None
+        self.diarization_version = diarization_version
 
     @staticmethod
     def has_libraries():
@@ -47,17 +48,17 @@ class Diarization:
 
     def initialize(self):
         """
-        1.Install pyannote.audio 3.0 with pip install pyannote.audio
+        1.Install pyannote.audio 3.1 with pip install pyannote.audio
         2.Accept pyannote/segmentation-3.0 user conditions
-        3.Accept pyannote/speaker-diarization-3.0 user conditions
+        3.Accept pyannote/speaker-diarization-3.1 user conditions
         4.Create access token at hf.co/settings/tokens.
-        https://huggingface.co/pyannote/speaker-diarization-3.0
+        https://huggingface.co/pyannote/speaker-diarization-3.1
         """
         if self.initialized:
             return
         from pyannote.audio import Pipeline
-
-        self.pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.0", use_auth_token=self.auth_token)
+        
+        self.pipeline = Pipeline.from_pretrained(f"pyannote/{self.diarization_version}", use_auth_token=self.auth_token)
         self.initialized = True
 
         # Load GPU mode if available
@@ -174,7 +175,7 @@ def main():
     # Read whisper JSON or SRT file
     whisper_result = load_transcript(args.whisper_file)
 
-    diarization = Diarization(auth_token=args.auth_token)
+    diarization = Diarization(auth_token=args.auth_token, diarization_version=args.diarization_version)
     diarization_result = list(diarization.run(args.audio_file, num_speakers=args.num_speakers, min_speakers=args.min_speakers, max_speakers=args.max_speakers))
 
     # Print result

src/diarization/diarizationContainer.py

@@ -4,13 +4,14 @@ from src.modelCache import GLOBAL_MODEL_CACHE, ModelCache
 from src.vadParallel import ParallelContext
 
 class DiarizationContainer:
-    def __init__(self, auth_token: str = None, enable_daemon_process: bool = True, auto_cleanup_timeout_seconds=60, cache: ModelCache = None):
+    def __init__(self, auth_token: str = None, enable_daemon_process: bool = True, auto_cleanup_timeout_seconds=60, cache: ModelCache = None, diarization_version=None):
         self.auth_token = auth_token
         self.enable_daemon_process = enable_daemon_process
         self.auto_cleanup_timeout_seconds = auto_cleanup_timeout_seconds
         self.diarization_context: ParallelContext = None
         self.cache = cache
         self.model = None
+        self.diarization_version = diarization_version
 
     def run(self, audio_file, **kwargs):
         # Create parallel context if needed
@@ -37,18 +38,18 @@ class DiarizationContainer:
             return self.model.mark_speakers(diarization_result, whisper_result)
 
         # Create a new diarization model (calling mark_speakers will not initialize pyannote.audio)
-        model = Diarization(self.auth_token)
+        model = Diarization(self.auth_token, self.diarization_version)
         return model.mark_speakers(diarization_result, whisper_result)
 
     def get_model(self):
         # Lazy load the model
         if (self.model is None):
             if self.cache:
-                print("Loading diarization model from cache")
-                self.model = self.cache.get("diarization", lambda : Diarization(self.auth_token))
+                print(f"Loading {self.diarization_version} model from cache")
+                self.model = self.cache.get(self.diarization_version, lambda : Diarization(self.auth_token, self.diarization_version))
             else:
-                print("Loading diarization model")
-                self.model = Diarization(self.auth_token)
+                print(f"Loading {self.diarization_version} model")
+                self.model = Diarization(self.auth_token, self.diarization_version)
         return self.model
 
     def execute(self, audio_file, **kwargs):
@@ -66,7 +67,8 @@ class DiarizationContainer:
         return {
             "auth_token": self.auth_token,
             "enable_daemon_process": self.enable_daemon_process,
-            "auto_cleanup_timeout_seconds": self.auto_cleanup_timeout_seconds
+            "auto_cleanup_timeout_seconds": self.auto_cleanup_timeout_seconds,
+            "diarization_version": self.diarization_version
         }
     
     def __setstate__(self, state):
@@ -74,5 +76,6 @@ class DiarizationContainer:
         self.enable_daemon_process = state["enable_daemon_process"]
         self.auto_cleanup_timeout_seconds = state["auto_cleanup_timeout_seconds"]
         self.diarization_context = None
+        self.diarization_version = state["diarization_version"]
         self.cache = GLOBAL_MODEL_CACHE
         self.model = None

src/utils.py

@@ -150,7 +150,7 @@ def __subtitle_preprocessor_iterator(transcript: Iterator[dict], maxLineWidth: i
            yield segment
 
         if segment_longest_speaker is not None:
-            segment_longest_speaker = segment_longest_speaker.replace("SPEAKER", "S")
+            segment_longest_speaker = "(" + segment_longest_speaker.replace("SPEAKER", "S") + ")"
             
         subtitle_start = segment['start']
         subtitle_end   = segment['end']
@@ -160,7 +160,9 @@ def __subtitle_preprocessor_iterator(transcript: Iterator[dict], maxLineWidth: i
         if len(words) == 0:
             # Prepend the longest speaker ID if available
             if segment_longest_speaker is not None:
-                text = f"({segment_longest_speaker}) {text}"
+                text = f"{segment_longest_speaker} {text}"
+                if text_original is not None and len(text_original) > 0:
+                    text_original = f"{segment_longest_speaker} {text_original}"
                 
             result = {
                 'start': subtitle_start,
@@ -175,12 +177,16 @@ def __subtitle_preprocessor_iterator(transcript: Iterator[dict], maxLineWidth: i
             continue
 
         if segment_longest_speaker is not None:
-            # Add the beginning
-            words.insert(0, {
-                'start': subtitle_start,
-                'end'  : subtitle_start,
-                'word' : f"({segment_longest_speaker})"
-            })
+            if words[0].get('word') != segment_longest_speaker:
+                # Add the beginning
+                words.insert(0, {
+                    'start': subtitle_start,
+                    'end'  : subtitle_start,
+                    'word' : segment_longest_speaker
+                })
+            text = f"{segment_longest_speaker} {text}"
+            if text_original is not None and len(text_original) > 0:
+                text_original = f"{segment_longest_speaker} {text_original}"
             
         text_words = [text] if not highlight_words and text_original is not None and len(text_original) > 0 else [ this_word["word"] for this_word in words ]