Merge branch 'main' of https://huggingface.co/spaces/aadnk/whisper-webui

Changes:

* Using pyannote/speaker-diarization-3.0.
* Display diarization min speakers and diarization max speakers options in the simple tab.
* When diarization speakers are set to 0, initialize diarization using only the min and max speakers options.

README.md

@@ -72,7 +72,7 @@ pip install -r requirements-fasterWhisper.txt
 ```
 And then run the App or the CLI with the `--whisper_implementation faster-whisper` flag:
 ```
-python app.py --whisper_implementation faster-whisper --input_audio_max_duration -1 --server_name 127.0.0.1 --auto_parallel True
+python app.py --whisper_implementation faster-whisper --input_audio_max_duration -1 --server_name 127.0.0.1 --server_port 7860 --auto_parallel True
 ```
 You can also select the whisper implementation in `config.json5`:
 ```json5

app.py

@@ -7,6 +7,7 @@ import argparse
 from io import StringIO
 import time
 import os
+import pathlib
 import tempfile
 import zipfile
 import numpy as np
@@ -14,6 +15,8 @@ import numpy as np
 import torch
 
 from src.config import VAD_INITIAL_PROMPT_MODE_VALUES, ApplicationConfig, VadInitialPromptMode
+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
@@ -33,7 +36,7 @@ import ffmpeg
 import gradio as gr
 
 from src.download import ExceededMaximumDuration, download_url
-from src.utils import optional_int, slugify, write_srt, write_vtt
+from src.utils import optional_int, slugify, str2bool, write_srt, write_vtt
 from src.vad import AbstractTranscription, NonSpeechStrategy, PeriodicTranscriptionConfig, TranscriptionConfig, VadPeriodicTranscription, VadSileroTranscription
 from src.whisper.abstractWhisperContainer import AbstractWhisperContainer
 from src.whisper.whisperFactory import create_whisper_container
@@ -83,6 +86,10 @@ class WhisperTranscriber:
         self.deleteUploadedFiles = delete_uploaded_files
         self.output_dir = output_dir
 
+        # Support for diarization
+        self.diarization: DiarizationContainer = None
+        # Dictionary with parameters to pass to diarization.run - if None, diarization is not enabled
+        self.diarization_kwargs = None
         self.app_config = app_config
 
     def set_parallel_devices(self, vad_parallel_devices: str):
@@ -96,22 +103,49 @@ 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):
+        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)
+        # Set parameters
+        self.diarization_kwargs = kwargs
+
+    def unset_diarization(self):
+        if self.diarization is not None:
+            self.diarization.cleanup()
+        self.diarization_kwargs = None
+
     # Entry function for the simple tab
     def transcribe_webui_simple(self, modelName, languageName, nllbModelName, nllbLangName, urlData, multipleFiles, microphoneData, task, 
                                 vad, vadMergeWindow, vadMaxMergeSize, 
-                                word_timestamps: bool = False, highlight_words: bool = False):
+                                word_timestamps: bool = False, highlight_words: bool = False,
+                                diarization: bool = False, diarization_speakers: int = 2,
+                                diarization_min_speakers = 1, diarization_max_speakers = 8):
         return self.transcribe_webui_simple_progress(modelName, languageName, nllbModelName, nllbLangName, urlData, multipleFiles, microphoneData, task, 
-                                                     vad, vadMergeWindow, vadMaxMergeSize, 
-                                                     word_timestamps, highlight_words)
+                                vad, vadMergeWindow, vadMaxMergeSize, 
+                                word_timestamps, highlight_words,
+                                diarization, diarization_speakers,
+                                diarization_min_speakers, diarization_max_speakers)
     
     # Entry function for the simple tab progress
     def transcribe_webui_simple_progress(self, modelName, languageName, nllbModelName, nllbLangName, urlData, multipleFiles, microphoneData, task, 
                                          vad, vadMergeWindow, vadMaxMergeSize, 
                                          word_timestamps: bool = False, highlight_words: bool = False, 
+                                         diarization: bool = False, diarization_speakers: int = 2,
+                                        diarization_min_speakers = 1, diarization_max_speakers = 8,
                                          progress=gr.Progress()):
-
+        
         vadOptions = VadOptions(vad, vadMergeWindow, vadMaxMergeSize, self.app_config.vad_padding, self.app_config.vad_prompt_window, self.app_config.vad_initial_prompt_mode)
 
+        if diarization:
+            if diarization_speakers < 1:
+                self.set_diarization(auth_token=self.app_config.auth_token, min_speakers=diarization_min_speakers, max_speakers=diarization_max_speakers)
+            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)
+        else:
+            self.unset_diarization()
+
         return self.transcribe_webui(modelName, languageName, nllbModelName, nllbLangName, urlData, multipleFiles, microphoneData, task, vadOptions, 
                                      word_timestamps=word_timestamps, highlight_words=highlight_words, progress=progress)
 
@@ -122,14 +156,18 @@ class WhisperTranscriber:
                               word_timestamps: bool, highlight_words: bool, prepend_punctuations: str, append_punctuations: str,
                               initial_prompt: str, temperature: float, best_of: int, beam_size: int, patience: float, length_penalty: float, suppress_tokens: str, 
                               condition_on_previous_text: bool, fp16: bool, temperature_increment_on_fallback: float, 
-                              compression_ratio_threshold: float, logprob_threshold: float, no_speech_threshold: float):
+                              compression_ratio_threshold: float, logprob_threshold: float, no_speech_threshold: float,
+                              diarization: bool = False, diarization_speakers: int = 2,
+                              diarization_min_speakers = 1, diarization_max_speakers = 8):
         
         return self.transcribe_webui_full_progress(modelName, languageName, nllbModelName, nllbLangName, urlData, multipleFiles, microphoneData, task, 
                                 vad, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow, vadInitialPromptMode,
                                 word_timestamps, highlight_words, prepend_punctuations, append_punctuations,
                                 initial_prompt, temperature, best_of, beam_size, patience, length_penalty, suppress_tokens,
                                 condition_on_previous_text, fp16, temperature_increment_on_fallback,
-                                compression_ratio_threshold, logprob_threshold, no_speech_threshold)
+                                compression_ratio_threshold, logprob_threshold, no_speech_threshold,
+                                diarization, diarization_speakers,
+                                diarization_min_speakers, diarization_max_speakers)
 
     # Entry function for the full tab with progress
     def transcribe_webui_full_progress(self, modelName, languageName, nllbModelName, nllbLangName, urlData, multipleFiles, microphoneData, task, 
@@ -139,6 +177,8 @@ class WhisperTranscriber:
                                         initial_prompt: str, temperature: float, best_of: int, beam_size: int, patience: float, length_penalty: float, suppress_tokens: str, 
                                         condition_on_previous_text: bool, fp16: bool, temperature_increment_on_fallback: float, 
                                         compression_ratio_threshold: float, logprob_threshold: float, no_speech_threshold: float, 
+                                        diarization: bool = False, diarization_speakers: int = 2, 
+                                        diarization_min_speakers = 1, diarization_max_speakers = 8,
                                         progress=gr.Progress()):
 
         # Handle temperature_increment_on_fallback
@@ -149,6 +189,15 @@ class WhisperTranscriber:
 
         vadOptions = VadOptions(vad, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow, vadInitialPromptMode)
 
+        # Set diarization
+        if diarization:
+            if diarization_speakers < 1:
+                self.set_diarization(auth_token=self.app_config.auth_token, min_speakers=diarization_min_speakers, max_speakers=diarization_max_speakers)
+            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)
+        else:
+            self.unset_diarization()
+
         return self.transcribe_webui(modelName, languageName, nllbModelName, nllbLangName, urlData, multipleFiles, microphoneData, task, vadOptions,
                                      initial_prompt=initial_prompt, temperature=temperature, best_of=best_of, beam_size=beam_size, patience=patience, length_penalty=length_penalty, suppress_tokens=suppress_tokens,
                                      condition_on_previous_text=condition_on_previous_text, fp16=fp16,
@@ -373,6 +422,19 @@ class WhisperTranscriber:
             else:
                 # Default VAD
                 result = whisperCallable.invoke(audio_path, 0, None, None, progress_listener=progressListener)
+        
+        # Diarization
+        if self.diarization and self.diarization_kwargs:
+            print("Diarizing ", audio_path)
+            diarization_result = list(self.diarization.run(audio_path, **self.diarization_kwargs))
+
+            # Print result
+            print("Diarization result: ")
+            for entry in diarization_result:
+                print(f"  start={entry.start:.1f}s stop={entry.end:.1f}s speaker_{entry.speaker}")
+
+            # Add speakers to result
+            result = self.diarization.mark_speakers(diarization_result, result)
 
         return result
 
@@ -545,11 +607,15 @@ class WhisperTranscriber:
         if (self.cpu_parallel_context is not None):
             self.cpu_parallel_context.close()
 
+        # Cleanup diarization
+        if (self.diarization is not None):
+            self.diarization.cleanup()
+            self.diarization = None
 
 def create_ui(app_config: ApplicationConfig):
     ui = WhisperTranscriber(app_config.input_audio_max_duration, app_config.vad_process_timeout, app_config.vad_cpu_cores, 
                             app_config.delete_uploaded_files, app_config.output_dir, app_config)
-    
+
     # Specify a list of devices to use for parallel processing
     ui.set_parallel_devices(app_config.vad_parallel_devices)
     ui.set_auto_parallel(app_config.auto_parallel)
@@ -619,6 +685,19 @@ def create_ui(app_config: ApplicationConfig):
         gr.Checkbox(label="Word Timestamps - Highlight Words", value=app_config.highlight_words),
     ]
 
+    has_diarization_libs = Diarization.has_libraries()
+
+    if not has_diarization_libs:
+        print("Diarization libraries not found - disabling diarization")
+        app_config.diarization = False
+
+    common_diarization_inputs = lambda : [
+        gr.Checkbox(label="Diarization", value=app_config.diarization, interactive=has_diarization_libs),
+        gr.Number(label="Diarization - Speakers", precision=0, value=app_config.diarization_speakers, interactive=has_diarization_libs),
+        gr.Number(label="Diarization - Min Speakers", precision=0, value=app_config.diarization_min_speakers, interactive=has_diarization_libs),
+        gr.Number(label="Diarization - Max Speakers", precision=0, value=app_config.diarization_max_speakers, interactive=has_diarization_libs)
+    ]
+    
     common_output = lambda : [
         gr.File(label="Download"),
         gr.Text(label="Transcription"),
@@ -640,7 +719,7 @@ def create_ui(app_config: ApplicationConfig):
                     with gr.Row():
                         simple_input += common_nllb_inputs()
                 with gr.Column():
-                    simple_input += common_audio_inputs() + common_vad_inputs() + common_word_timestamps_inputs()
+                    simple_input += common_audio_inputs() + common_vad_inputs() + common_word_timestamps_inputs() + common_diarization_inputs()
             with gr.Column():
                 simple_output = common_output()
                 simple_flag = gr.Button("Flag")
@@ -689,7 +768,7 @@ def create_ui(app_config: ApplicationConfig):
                     gr.Number(label="Temperature increment on fallback", value=app_config.temperature_increment_on_fallback),
                     gr.Number(label="Compression ratio threshold", value=app_config.compression_ratio_threshold),
                     gr.Number(label="Logprob threshold", value=app_config.logprob_threshold),
-                    gr.Number(label="No speech threshold", value=app_config.no_speech_threshold)]
+                    gr.Number(label="No speech threshold", value=app_config.no_speech_threshold)] + common_diarization_inputs()
 
             with gr.Column():
                 full_output = common_output()
@@ -771,7 +850,14 @@ if __name__ == '__main__':
                         help="Maximum length of a file name.")
     parser.add_argument("--autolaunch", action='store_true', \
                         help="open the webui URL in the system's default browser upon launch")
-                        
+    parser.add_argument('--auth_token', type=str, default=default_app_config.auth_token, help='HuggingFace API Token (optional)')
+    parser.add_argument("--diarization", type=str2bool, default=default_app_config.diarization, \
+                        help="whether to perform speaker diarization")
+    parser.add_argument("--diarization_num_speakers", type=int, default=default_app_config.diarization_speakers, help="Number of speakers")
+    parser.add_argument("--diarization_min_speakers", type=int, default=default_app_config.diarization_min_speakers, help="Minimum number of speakers")
+    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.")
 
     args = parser.parse_args().__dict__
 
@@ -788,4 +874,5 @@ if __name__ == '__main__':
     if (threads := args.pop("threads")) > 0:
         torch.set_num_threads(threads)
 
+    print("Using whisper implementation: " + updated_config.whisper_implementation)
     create_ui(app_config=updated_config)

cli.py

@@ -8,6 +8,7 @@ import numpy as np
 import torch
 from app import VadOptions, WhisperTranscriber
 from src.config import VAD_INITIAL_PROMPT_MODE_VALUES, ApplicationConfig, VadInitialPromptMode
+from src.diarization.diarization import Diarization
 from src.download import download_url
 from src.languages import get_language_names
 
@@ -106,6 +107,14 @@ def cli():
     parser.add_argument("--threads", type=optional_int, default=0, 
                         help="number of threads used by torch for CPU inference; supercedes MKL_NUM_THREADS/OMP_NUM_THREADS")
 
+    # Diarization
+    parser.add_argument('--auth_token', type=str, default=None, help='HuggingFace API Token (optional)')
+    parser.add_argument("--diarization", type=str2bool, default=app_config.diarization, \
+                        help="whether to perform speaker diarization")
+    parser.add_argument("--diarization_num_speakers", type=int, default=None, help="Number of speakers")
+    parser.add_argument("--diarization_min_speakers", type=int, default=None, help="Minimum number of speakers")
+    parser.add_argument("--diarization_max_speakers", type=int, default=None, help="Maximum number of speakers")
+
     args = parser.parse_args().__dict__
     model_name: str = args.pop("model")
     model_dir: str = args.pop("model_dir")
@@ -142,10 +151,19 @@ def cli():
     compute_type = args.pop("compute_type")
     highlight_words = args.pop("highlight_words")
 
+    auth_token = args.pop("auth_token")
+    diarization = args.pop("diarization")
+    num_speakers = args.pop("diarization_num_speakers")
+    min_speakers = args.pop("diarization_min_speakers")
+    max_speakers = args.pop("diarization_max_speakers")
+    
     transcriber = WhisperTranscriber(delete_uploaded_files=False, vad_cpu_cores=vad_cpu_cores, app_config=app_config)
     transcriber.set_parallel_devices(args.pop("vad_parallel_devices"))
     transcriber.set_auto_parallel(auto_parallel)
 
+    if diarization:
+        transcriber.set_diarization(auth_token=auth_token, enable_daemon_process=False, num_speakers=num_speakers, min_speakers=min_speakers, max_speakers=max_speakers)
+
     model = create_whisper_container(whisper_implementation=whisper_implementation, model_name=model_name, 
                                      device=device, compute_type=compute_type, download_root=model_dir, models=app_config.models)
     

config.json5

@@ -234,4 +234,17 @@
     "append_punctuations": "\"\'.。,，!！?？:：”)]}、",
     // (requires --word_timestamps True) underline each word as it is spoken in srt and vtt
     "highlight_words": false,
+
+    // Diarization settings
+    "auth_token": null,
+    // Whether to perform speaker diarization
+    "diarization": false,
+    // The number of speakers to detect
+    "diarization_speakers": 2,
+    // The minimum number of speakers to detect
+    "diarization_min_speakers": 1,
+    // The maximum number of speakers to detect
+    "diarization_max_speakers": 8,
+    // The number of seconds before inactivate processes are terminated. Use 0 to close processes immediately, or None for no timeout.
+    "diarization_process_timeout": 60,
 }

docs/options.md

@@ -80,6 +80,17 @@ number of seconds after the line has finished. For instance, if a line ends at 1
 Note that detected lines in gaps between speech sections will not be included in the prompt 
 (if silero-vad or silero-vad-expand-into-gaps) is used.
 
+## Diarization
+
+If checked, Pyannote will be used to detect speakers in the audio, and label them as (SPEAKER 00), (SPEAKER 01), etc. 
+
+This requires a HuggingFace API key to function, which can be supplied with the `--auth_token` command line option for the CLI,
+set in the `config.json5` file for the GUI, or provided via the `HF_ACCESS_TOKEN` environment variable.
+
+## Diarization - Speakers
+
+The number of speakers to detect. If set to 0, Pyannote will attempt to detect the number of speakers automatically.
+
 # Command Line Options
 
 Both `app.py` and `cli.py` also accept command line options, such as the ability to enable parallel execution on multiple
@@ -132,3 +143,11 @@ If the average log probability is lower than this value, treat the decoding as f
 
 ## No speech threshold
 If the probability of the <|nospeech|> token is higher than this value AND the decoding has failed due to `logprob_threshold`, consider the segment as silence. Default is 0.6.
+
+## Diarization - Min Speakers
+
+The minimum number of speakers for Pyannote to detect.
+
+## Diarization - Max Speakers
+
+The maximum number of speakers for Pyannote to detect.

requirements-fasterWhisper.txt

@@ -9,4 +9,10 @@ torch
 torchaudio
 more_itertools
 zhconv
-sentencepiece
+sentencepiece
+
+# Needed by diarization
+intervaltree
+srt
+torch
+https://github.com/pyannote/pyannote-audio/archive/refs/heads/develop.zip

requirements-whisper.txt

@@ -8,4 +8,10 @@ torchaudio
 altair
 json5
 zhconv
-sentencepiece
+sentencepiece
+
+# Needed by diarization
+intervaltree
+srt
+torch
+https://github.com/pyannote/pyannote-audio/archive/refs/heads/develop.zip

requirements.txt

@@ -9,4 +9,10 @@ torch
 torchaudio
 more_itertools
 zhconv
-sentencepiece
+sentencepiece
+
+# Needed by diarization
+intervaltree
+srt
+torch
+https://github.com/pyannote/pyannote-audio/archive/refs/heads/develop.zip

src/config.py

@@ -69,7 +69,11 @@ class ApplicationConfig:
                  # Word timestamp settings
                  word_timestamps: bool = False, prepend_punctuations: str = "\"\'“¿([{-",
                  append_punctuations: str = "\"\'.。,，!！?？:：”)]}、", 
-                 highlight_words: bool = False):
+                 highlight_words: bool = False,
+                 # Diarization
+                 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):
         
         self.models = models
         self.nllb_models = nllb_models
@@ -123,6 +127,14 @@ class ApplicationConfig:
         self.append_punctuations = append_punctuations
         self.highlight_words = highlight_words
         
+        # Diarization settings
+        self.auth_token = auth_token
+        self.diarization = diarization
+        self.diarization_speakers = diarization_speakers
+        self.diarization_min_speakers = diarization_min_speakers
+        self.diarization_max_speakers = diarization_max_speakers
+        self.diarization_process_timeout = diarization_process_timeout
+
     def get_model_names(self):
         return [ x.name for x in self.models ]
 

src/diarization/diarization.py

@@ -0,0 +1,202 @@
+import argparse
+import gc
+import json
+import os
+from pathlib import Path
+import tempfile
+from typing import TYPE_CHECKING, List
+import torch
+
+import ffmpeg
+
+class DiarizationEntry:
+    def __init__(self, start, end, speaker):
+        self.start = start
+        self.end = end
+        self.speaker = speaker
+
+    def __repr__(self):
+        return f"<DiarizationEntry start={self.start} end={self.end} speaker={self.speaker}>"
+    
+    def toJson(self):
+        return {
+            "start": self.start,
+            "end": self.end,
+            "speaker": self.speaker
+        }
+
+class Diarization:
+    def __init__(self, auth_token=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
+
+    @staticmethod
+    def has_libraries():
+        try:
+            import pyannote.audio
+            import intervaltree
+            return True
+        except ImportError:
+            return False
+
+    def initialize(self):
+        """
+        1.Install pyannote.audio 3.0 with pip install pyannote.audio
+        2.Accept pyannote/segmentation-3.0 user conditions
+        3.Accept pyannote/speaker-diarization-3.0 user conditions
+        4.Create access token at hf.co/settings/tokens.
+        https://huggingface.co/pyannote/speaker-diarization-3.0
+        """
+        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.initialized = True
+
+        # Load GPU mode if available
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        if device == "cuda":
+            print("Diarization - using GPU")
+            self.pipeline = self.pipeline.to(torch.device(0))
+        else:
+            print("Diarization - using CPU")
+
+    def run(self, audio_file, **kwargs):
+        self.initialize()
+        audio_file_obj = Path(audio_file)
+
+        # Supported file types in soundfile is WAV, FLAC, OGG and MAT
+        if audio_file_obj.suffix in [".wav", ".flac", ".ogg", ".mat"]:
+            target_file = audio_file
+        else:
+            # Create temp WAV file
+            target_file = tempfile.mktemp(prefix="diarization_", suffix=".wav")
+            try:
+                ffmpeg.input(audio_file).output(target_file, ac=1).run()
+            except ffmpeg.Error as e:
+                print(f"Error occurred during audio conversion: {e.stderr}")
+
+        diarization = self.pipeline(target_file, **kwargs)
+
+        if target_file != audio_file:
+            # Delete temp file
+            os.remove(target_file)
+
+        # Yield result
+        for turn, _, speaker in diarization.itertracks(yield_label=True):
+            yield DiarizationEntry(turn.start, turn.end, speaker)
+    
+    def mark_speakers(self, diarization_result: List[DiarizationEntry], whisper_result: dict):
+        from intervaltree import IntervalTree
+        result = whisper_result.copy()
+
+        # Create an interval tree from the diarization results
+        tree = IntervalTree()
+        for entry in diarization_result:
+            tree[entry.start:entry.end] = entry
+
+        # Iterate through each segment in the Whisper JSON
+        for segment in result["segments"]:
+            segment_start = segment["start"]
+            segment_end = segment["end"]
+
+            # Find overlapping speakers using the interval tree
+            overlapping_speakers = tree[segment_start:segment_end]
+
+            # If no speakers overlap with this segment, skip it
+            if not overlapping_speakers:
+                continue
+
+            # If multiple speakers overlap with this segment, choose the one with the longest duration
+            longest_speaker = None
+            longest_duration = 0
+            
+            for speaker_interval in overlapping_speakers:
+                overlap_start = max(speaker_interval.begin, segment_start)
+                overlap_end = min(speaker_interval.end, segment_end)
+                overlap_duration = overlap_end - overlap_start
+
+                if overlap_duration > longest_duration:
+                    longest_speaker = speaker_interval.data.speaker
+                    longest_duration = overlap_duration
+
+            # Add speakers
+            segment["longest_speaker"] = longest_speaker
+            segment["speakers"] = list([speaker_interval.data.toJson() for speaker_interval in overlapping_speakers])
+
+            # The write_srt will use the longest_speaker if it exist, and add it to the text field
+
+        return result
+
+def _write_file(input_file: str, output_path: str, output_extension: str, file_writer: lambda f: None):
+    if input_file is None:
+        raise ValueError("input_file is required")
+    if file_writer is None:
+        raise ValueError("file_writer is required")
+
+     # Write file
+    if output_path is None:
+        effective_path = os.path.splitext(input_file)[0] + "_output" + output_extension
+    else:
+        effective_path = output_path
+
+    with open(effective_path, 'w+', encoding="utf-8") as f:
+        file_writer(f)
+
+    print(f"Output saved to {effective_path}")
+
+def main():
+    from src.utils import write_srt
+    from src.diarization.transcriptLoader import load_transcript
+
+    parser = argparse.ArgumentParser(description='Add speakers to a SRT file or Whisper JSON file using pyannote/speaker-diarization.')
+    parser.add_argument('audio_file', type=str, help='Input audio file')
+    parser.add_argument('whisper_file', type=str, help='Input Whisper JSON/SRT file')
+    parser.add_argument('--output_json_file', type=str, default=None, help='Output JSON file (optional)')
+    parser.add_argument('--output_srt_file', type=str, default=None, help='Output SRT file (optional)')
+    parser.add_argument('--auth_token', type=str, default=None, help='HuggingFace API Token (optional)')
+    parser.add_argument("--max_line_width", type=int, default=40, help="Maximum line width for SRT file (default: 40)")
+    parser.add_argument("--num_speakers", type=int, default=None, help="Number of speakers")
+    parser.add_argument("--min_speakers", type=int, default=None, help="Minimum number of speakers")
+    parser.add_argument("--max_speakers", type=int, default=None, help="Maximum number of speakers")
+
+    args = parser.parse_args()
+
+    print("\nReading whisper JSON from " + args.whisper_file)
+
+    # Read whisper JSON or SRT file
+    whisper_result = load_transcript(args.whisper_file)
+
+    diarization = Diarization(auth_token=args.auth_token)
+    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
+    print("Diarization result:")
+    for entry in diarization_result:
+        print(f"  start={entry.start:.1f}s stop={entry.end:.1f}s speaker_{entry.speaker}")
+
+    marked_whisper_result = diarization.mark_speakers(diarization_result, whisper_result)
+
+    # Write output JSON to file
+    _write_file(args.whisper_file, args.output_json_file, ".json", 
+                lambda f: json.dump(marked_whisper_result, f, indent=4, ensure_ascii=False))
+
+    # Write SRT
+    _write_file(args.whisper_file, args.output_srt_file, ".srt", 
+                lambda f: write_srt(marked_whisper_result["segments"], f, maxLineWidth=args.max_line_width))
+
+if __name__ == "__main__":
+    main()
+    
+    #test = Diarization()
+    #print("Initializing")
+    #test.initialize()
+
+    #input("Press Enter to continue...")

src/diarization/diarizationContainer.py

@@ -0,0 +1,78 @@
+from typing import List
+from src.diarization.diarization import Diarization, DiarizationEntry
+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):
+        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
+
+    def run(self, audio_file, **kwargs):
+        # Create parallel context if needed
+        if self.diarization_context is None and self.enable_daemon_process:
+            # Number of processes is set to 1 as we mainly use this in order to clean up GPU memory
+            self.diarization_context = ParallelContext(num_processes=1, auto_cleanup_timeout_seconds=self.auto_cleanup_timeout_seconds)
+            print("Created diarization context with auto cleanup timeout of %d seconds" % self.auto_cleanup_timeout_seconds)
+        
+        # Run directly 
+        if self.diarization_context is None:
+            return self.execute(audio_file, **kwargs)
+
+        # Otherwise run in a separate process
+        pool = self.diarization_context.get_pool()
+
+        try:
+            result = pool.apply(self.execute, (audio_file,), kwargs)
+            return result
+        finally:
+            self.diarization_context.return_pool(pool)
+
+    def mark_speakers(self, diarization_result: List[DiarizationEntry], whisper_result: dict):
+        if self.model is not None:
+            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)
+        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))
+            else:
+                print("Loading diarization model")
+                self.model = Diarization(self.auth_token)
+        return self.model
+
+    def execute(self, audio_file, **kwargs):
+        model = self.get_model()
+
+        # We must use list() here to force the iterator to run, as generators are not picklable
+        result = list(model.run(audio_file, **kwargs))
+        return result
+    
+    def cleanup(self):
+        if self.diarization_context is not None:
+            self.diarization_context.close()
+
+    def __getstate__(self):
+        return {
+            "auth_token": self.auth_token,
+            "enable_daemon_process": self.enable_daemon_process,
+            "auto_cleanup_timeout_seconds": self.auto_cleanup_timeout_seconds
+        }
+    
+    def __setstate__(self, state):
+        self.auth_token = state["auth_token"]
+        self.enable_daemon_process = state["enable_daemon_process"]
+        self.auto_cleanup_timeout_seconds = state["auto_cleanup_timeout_seconds"]
+        self.diarization_context = None
+        self.cache = GLOBAL_MODEL_CACHE
+        self.model = None

src/diarization/requirements.txt

@@ -0,0 +1,5 @@
+intervaltree
+srt
+torch
+ffmpeg-python==0.2.0
+https://github.com/pyannote/pyannote-audio/archive/refs/heads/develop.zip

src/diarization/transcriptLoader.py

@@ -0,0 +1,80 @@
+
+import json
+from pathlib import Path
+
+def load_transcript_json(transcript_file: str):
+    """
+    Parse a Whisper JSON file into a Whisper JSON object
+
+    # Parameters:
+    transcript_file (str): Path to the Whisper JSON file
+    """
+    with open(transcript_file, "r", encoding="utf-8") as f:
+        whisper_result = json.load(f)
+
+    # Format of Whisper JSON file:
+    #  {
+    # "text": " And so my fellow Americans, ask not what your country can do for you, ask what you can do for your country.",
+    # "segments": [
+    #    {
+    #        "text": " And so my fellow Americans, ask not what your country can do for you, ask what you can do for your country.",
+    #        "start": 0.0,
+    #        "end": 10.36,
+    #        "words": [
+    #            {
+    #                "start": 0.0,
+    #                "end": 0.56,
+    #                "word": " And",
+    #                "probability": 0.61767578125
+    #            },
+    #            {
+    #                "start": 0.56,
+    #                "end": 0.88,
+    #                "word": " so",
+    #                "probability": 0.9033203125
+    #            },
+    # etc.  
+
+    return whisper_result
+
+
+def load_transcript_srt(subtitle_file: str):
+    import srt
+
+    """
+    Parse a SRT file into a Whisper JSON object
+
+    # Parameters:
+    subtitle_file (str): Path to the SRT file
+    """
+    with open(subtitle_file, "r", encoding="utf-8") as f:
+        subs = srt.parse(f)
+
+        whisper_result = {
+            "text": "",
+            "segments": []
+        }
+
+        for sub in subs:
+            # Subtitle(index=1, start=datetime.timedelta(seconds=33, microseconds=843000), end=datetime.timedelta(seconds=38, microseconds=97000), content='地球上只有3%的水是淡水', proprietary='')
+            segment = {
+                "text": sub.content,
+                "start": sub.start.total_seconds(),
+                "end": sub.end.total_seconds(),
+                "words": []
+            }
+            whisper_result["segments"].append(segment)
+            whisper_result["text"] += sub.content
+
+        return whisper_result
+
+def load_transcript(file: str):
+    # Determine file type
+    file_extension = Path(file).suffix.lower()
+
+    if file_extension == ".json":
+        return load_transcript_json(file)
+    elif file_extension == ".srt":
+        return load_transcript_srt(file)
+    else:
+        raise ValueError(f"Unsupported file type: {file_extension}")

src/utils.py

@@ -102,17 +102,28 @@ def write_srt(transcript: Iterator[dict], file: TextIO,
 
 def __subtitle_preprocessor_iterator(transcript: Iterator[dict], maxLineWidth: int = None, highlight_words: bool = False): 
     for segment in transcript:
-        words = segment.get('words', [])
+        words: list = segment.get('words', [])
+
+        # Append longest speaker ID if available
+        segment_longest_speaker = segment.get('longest_speaker', None)
+        if segment_longest_speaker is not None:
+            segment_longest_speaker = segment_longest_speaker.replace("SPEAKER", "S")
 
         if len(words) == 0:
             # Yield the segment as-is or processed
-            if maxLineWidth is None or maxLineWidth < 0:
+            if (maxLineWidth is None or maxLineWidth < 0) and segment_longest_speaker is None:
                 yield segment
             else:
+                text = segment['text'].strip()
+
+                # Prepend the longest speaker ID if available
+                if segment_longest_speaker is not None:
+                    text = f"({segment_longest_speaker}) {text}"
+
                 yield {
                     'start': segment['start'],
                     'end': segment['end'],
-                    'text': process_text(segment['text'].strip(), maxLineWidth)
+                    'text': process_text(text, maxLineWidth)
                 }
             # We are done
             continue
@@ -120,9 +131,17 @@ def __subtitle_preprocessor_iterator(transcript: Iterator[dict], maxLineWidth: i
         subtitle_start = segment['start']
         subtitle_end = segment['end']
 
+        if segment_longest_speaker is not None:
+            # Add the beginning
+            words.insert(0, {
+                'start': subtitle_start,
+                'end': subtitle_start,
+                'word': f"({segment_longest_speaker})"
+            })
+
         text_words = [ this_word["word"] for this_word in words ]
         subtitle_text = __join_words(text_words, maxLineWidth)
-        
+
         # Iterate over the words in the segment
         if highlight_words:
             last = subtitle_start

src/whisper/abstractWhisperContainer.py

@@ -1,5 +1,5 @@
 import abc
-from typing import List
+from typing import Any, Callable, List
 
 from src.config import ModelConfig, VadInitialPromptMode
 
@@ -9,7 +9,7 @@ from src.prompts.abstractPromptStrategy import AbstractPromptStrategy
 
 class AbstractWhisperCallback:
     def __init__(self):
-        self.__prompt_mode_gpt = None
+        pass
 
     @abc.abstractmethod
     def invoke(self, audio, segment_index: int, prompt: str, detected_language: str, progress_listener: ProgressListener = None):
@@ -29,6 +29,14 @@ class AbstractWhisperCallback:
         """
         raise NotImplementedError()
 
+class LambdaWhisperCallback(AbstractWhisperCallback):
+    def __init__(self, callback_lambda: Callable[[Any, int, str, str, ProgressListener], None]):
+        super().__init__()
+        self.callback_lambda = callback_lambda
+
+    def invoke(self, audio, segment_index: int, prompt: str, detected_language: str, progress_listener: ProgressListener = None):
+        return self.callback_lambda(audio, segment_index, prompt, detected_language, progress_listener)
+
 class AbstractWhisperContainer:
     def __init__(self, model_name: str, device: str = None, compute_type: str = "float16",
                  download_root: str = None,

src/whisper/dummyWhisperContainer.py

@@ -0,0 +1,101 @@
+from typing import List
+
+import ffmpeg
+from src.config import ModelConfig
+from src.hooks.progressListener import ProgressListener
+from src.modelCache import ModelCache
+from src.prompts.abstractPromptStrategy import AbstractPromptStrategy
+from src.whisper.abstractWhisperContainer import AbstractWhisperCallback, AbstractWhisperContainer
+
+class DummyWhisperContainer(AbstractWhisperContainer):
+    def __init__(self, model_name: str, device: str = None, compute_type: str = "float16",
+                       download_root: str = None,
+                       cache: ModelCache = None, models: List[ModelConfig] = []):
+        super().__init__(model_name, device, compute_type, download_root, cache, models)
+
+    def ensure_downloaded(self):
+        """
+        Ensure that the model is downloaded. This is useful if you want to ensure that the model is downloaded before
+        passing the container to a subprocess.
+        """
+        print("[Dummy] Ensuring that the model is downloaded")
+
+    def _create_model(self):
+        print("[Dummy] Creating dummy whisper model " + self.model_name + " for device " + str(self.device))
+        return None
+
+    def create_callback(self, language: str = None, task: str = None, 
+                        prompt_strategy: AbstractPromptStrategy = None, 
+                        **decodeOptions: dict) -> AbstractWhisperCallback:
+        """
+        Create a WhisperCallback object that can be used to transcript audio files.
+
+        Parameters
+        ----------
+        language: str
+            The target language of the transcription. If not specified, the language will be inferred from the audio content.
+        task: str
+            The task - either translate or transcribe.
+        prompt_strategy: AbstractPromptStrategy
+            The prompt strategy to use. If not specified, the prompt from Whisper will be used.
+        decodeOptions: dict
+            Additional options to pass to the decoder. Must be pickleable.
+
+        Returns
+        -------
+        A WhisperCallback object.
+        """
+        return DummyWhisperCallback(self, language=language, task=task, prompt_strategy=prompt_strategy, **decodeOptions)
+    
+class DummyWhisperCallback(AbstractWhisperCallback):
+    def __init__(self, model_container: DummyWhisperContainer, **decodeOptions: dict):
+        self.model_container = model_container
+        self.decodeOptions = decodeOptions
+        
+    def invoke(self, audio, segment_index: int, prompt: str, detected_language: str, progress_listener: ProgressListener = None):
+        """
+        Peform the transcription of the given audio file or data.
+
+        Parameters
+        ----------
+        audio: Union[str, np.ndarray, torch.Tensor]
+            The audio file to transcribe, or the audio data as a numpy array or torch tensor.
+        segment_index: int
+            The target language of the transcription. If not specified, the language will be inferred from the audio content.
+        task: str
+            The task - either translate or transcribe.
+        progress_listener: ProgressListener
+            A callback to receive progress updates.
+        """
+        print("[Dummy] Invoking dummy whisper callback for segment " + str(segment_index))
+
+        # Estimate length
+        if isinstance(audio, str):
+            audio_length = ffmpeg.probe(audio)["format"]["duration"]
+        # Format is pcm_s16le at a sample rate of 16000, loaded as a float32 array.
+        else:
+            audio_length = len(audio) / 16000
+
+        # Convert the segments to a format that is easier to serialize
+        whisper_segments = [{
+            "text": "Dummy text for segment " + str(segment_index),
+            "start": 0,
+            "end": audio_length,
+
+            # Extra fields added by faster-whisper
+            "words": []
+        }]
+
+        result = {
+            "segments": whisper_segments,
+            "text": "Dummy text for segment " + str(segment_index),
+            "language": "en" if detected_language is None else detected_language,
+
+            # Extra fields added by faster-whisper
+            "language_probability": 1.0,
+            "duration": audio_length,
+        }
+
+        if progress_listener is not None:
+            progress_listener.on_finished()
+        return result

src/whisper/whisperFactory.py

@@ -15,5 +15,9 @@ def create_whisper_container(whisper_implementation: str,
     elif (whisper_implementation == "faster-whisper" or whisper_implementation == "faster_whisper"):
         from src.whisper.fasterWhisperContainer import FasterWhisperContainer
         return FasterWhisperContainer(model_name=model_name, device=device, compute_type=compute_type, download_root=download_root, cache=cache, models=models)
+    elif (whisper_implementation == "dummy-whisper" or whisper_implementation == "dummy_whisper" or whisper_implementation == "dummy"):
+        # This is useful for testing
+        from src.whisper.dummyWhisperContainer import DummyWhisperContainer
+        return DummyWhisperContainer(model_name=model_name, device=device, compute_type=compute_type, download_root=download_root, cache=cache, models=models)
     else:
         raise ValueError("Unknown Whisper implementation: " + whisper_implementation)

tests/vad_test.py

@@ -1,10 +1,11 @@
-import pprint
 import unittest
 import numpy as np
 import sys
 
 sys.path.append('../whisper-webui')
+#print("Sys path: " + str(sys.path))
 
+from src.whisper.abstractWhisperContainer import LambdaWhisperCallback
 from src.vad import AbstractTranscription, TranscriptionConfig, VadSileroTranscription
 
 class TestVad(unittest.TestCase):
@@ -13,10 +14,11 @@ class TestVad(unittest.TestCase):
         self.transcribe_calls = []
 
     def test_transcript(self):
-        mock = MockVadTranscription()
+        mock = MockVadTranscription(mock_audio_length=120)
+        config = TranscriptionConfig()
 
         self.transcribe_calls.clear()
-        result = mock.transcribe("mock", lambda segment : self.transcribe_segments(segment))
+        result = mock.transcribe("mock", LambdaWhisperCallback(lambda segment, _1, _2, _3, _4: self.transcribe_segments(segment)), config)
 
         self.assertListEqual(self.transcribe_calls, [ 
             [30, 30],
@@ -45,8 +47,9 @@ class TestVad(unittest.TestCase):
         }
 
 class MockVadTranscription(AbstractTranscription):
-    def __init__(self):
+    def __init__(self, mock_audio_length: float = 1000):
         super().__init__()
+        self.mock_audio_length = mock_audio_length
 
     def get_audio_segment(self, str, start_time: str = None, duration: str = None):
         start_time_seconds = float(start_time.removesuffix("s"))
@@ -61,6 +64,9 @@ class MockVadTranscription(AbstractTranscription):
         result.append( {  'start': 30, 'end': 60 } )
         result.append( {  'start': 100, 'end': 200 } )
         return result
+        
+    def get_audio_duration(self, audio: str, config: TranscriptionConfig):
+        return self.mock_audio_length
 
 if __name__ == '__main__':
     unittest.main()