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

app.py

@@ -14,6 +14,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
@@ -73,6 +75,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):
@@ -86,22 +92,41 @@ 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.vad_process_timeout, cache=self.model_cache)
+        # Set parameters
+        self.diarization_kwargs = kwargs
+
+    def unset_diarization(self):
+        self.diarization.cleanup()
+        self.diarization_kwargs = None
+
     # Entry function for the simple tab
     def transcribe_webui_simple(self, modelName, languageName, 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):
         return self.transcribe_webui_simple_progress(modelName, languageName, urlData, multipleFiles, microphoneData, task, 
                                                      vad, vadMergeWindow, vadMaxMergeSize, 
-                                                     word_timestamps, highlight_words)
+                                                     word_timestamps, highlight_words, 
+                                                     diarization, diarization_speakers)
     
     # Entry function for the simple tab progress
     def transcribe_webui_simple_progress(self, modelName, languageName, urlData, multipleFiles, microphoneData, task, 
                                          vad, vadMergeWindow, vadMaxMergeSize, 
                                          word_timestamps: bool = False, highlight_words: bool = False, 
+                                         diarization: bool = False, diarization_speakers: int = 2,
                                          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:
+            self.set_diarization(auth_token=self.app_config.auth_token, num_speakers=diarization_speakers)
+        else:
+            self.unset_diarization()
+
         return self.transcribe_webui(modelName, languageName, urlData, multipleFiles, microphoneData, task, vadOptions, 
                                      word_timestamps=word_timestamps, highlight_words=highlight_words, progress=progress)
 
@@ -112,14 +137,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 = 5):
         
         return self.transcribe_webui_full_progress(modelName, languageName, 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, urlData, multipleFiles, microphoneData, task, 
@@ -129,6 +158,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 = 5,
                                         progress=gr.Progress()):
 
         # Handle temperature_increment_on_fallback
@@ -139,6 +170,13 @@ class WhisperTranscriber:
 
         vadOptions = VadOptions(vad, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow, vadInitialPromptMode)
 
+        # Set diarization
+        if diarization:
+            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, 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,
@@ -322,6 +360,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
 
@@ -458,6 +509,10 @@ 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, 
@@ -515,6 +570,17 @@ 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)
+    ]
+
     is_queue_mode = app_config.queue_concurrency_count is not None and app_config.queue_concurrency_count > 0    
 
     simple_transcribe = gr.Interface(fn=ui.transcribe_webui_simple_progress if is_queue_mode else ui.transcribe_webui_simple, 
@@ -522,6 +588,7 @@ def create_ui(app_config: ApplicationConfig):
         *common_inputs(),
         *common_vad_inputs(),
         *common_word_timestamps_inputs(),
+        *common_diarization_inputs(),
     ], outputs=[
         gr.File(label="Download"),
         gr.Text(label="Transcription"), 
@@ -556,6 +623,11 @@ def create_ui(app_config: ApplicationConfig):
         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),
+
+        *common_diarization_inputs(),
+        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),
+
     ], outputs=[
         gr.File(label="Download"),
         gr.Text(label="Transcription"), 

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

@@ -140,4 +140,15 @@
     "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": 5,
 }

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 `HK_AUTH_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

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

requirements-whisper.txt

@@ -6,4 +6,10 @@ gradio==3.38.0
 yt-dlp
 torchaudio
 altair
-json5
+json5
+
+# Needed by diarization
+intervaltree
+srt
+torch
+https://github.com/pyannote/pyannote-audio/archive/refs/heads/develop.zip

requirements.txt

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

src/config.py

@@ -69,7 +69,10 @@ 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):
         
         self.models = models
         
@@ -121,6 +124,13 @@ 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
+
     def get_model_names(self):
         return [ x.name for x in self.models ]
 

src/diarization/diarization.py

@@ -0,0 +1,195 @@
+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("HK_ACCESS_TOKEN")
+            if auth_token is None:
+                raise ValueError("No HuggingFace API Token provided - please use the --auth_token argument or set the HK_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):
+        if self.initialized:
+            return
+        from pyannote.audio import Pipeline
+
+        self.pipeline = Pipeline.from_pretrained("pyannote/[email protected]", 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,77 @@
+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)
+        
+        # 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,26 @@ 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 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 +129,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