Adding support for diarization

app.py

@@ -14,6 +14,7 @@ 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.hooks.progressListener import ProgressListener
 from src.hooks.subTaskProgressListener import SubTaskProgressListener
 from src.hooks.whisperProgressHook import create_progress_listener_handle
@@ -73,6 +74,7 @@ class WhisperTranscriber:
         self.deleteUploadedFiles = delete_uploaded_files
         self.output_dir = output_dir
 
+        self.diarization: Diarization = None
         self.app_config = app_config
 
     def set_parallel_devices(self, vad_parallel_devices: str):
@@ -89,19 +91,27 @@ class WhisperTranscriber:
     # 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.diarization = Diarization(auth_token=self.app_config.auth_token, num_speakers=diarization_speakers)
+        else:
+            self.diarization = None
+
         return self.transcribe_webui(modelName, languageName, urlData, multipleFiles, microphoneData, task, vadOptions, 
                                      word_timestamps=word_timestamps, highlight_words=highlight_words, progress=progress)
 
@@ -112,14 +122,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 +143,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 +155,13 @@ class WhisperTranscriber:
 
         vadOptions = VadOptions(vad, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow, vadInitialPromptMode)
 
+        # Set diarization
+        if diarization:
+            self.diarization = Diarization(auth_token=self.app_config.auth_token, num_speakers=diarization_speakers, 
+                                           min_speakers=diarization_min_speakers, max_speakers=diarization_max_speakers)
+        else:
+            self.diarization = None
+
         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,
@@ -202,6 +225,19 @@ class WhisperTranscriber:
                     # Update progress
                     current_progress += source_audio_duration
 
+                    # Diarization
+                    if self.diarization:
+                        print("Diarizing ", source.source_path)
+                        diarization_result = list(self.diarization.run(source.source_path))
+
+                        # 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)
+
                     source_download, source_text, source_vtt = self.write_result(result, filePrefix, outputDirectory, highlight_words)
 
                     if len(sources) > 1:
@@ -515,6 +551,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 +569,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 +604,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("--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().__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")
 
+    diarization = args.pop("diarization")
+    auth_token = args.pop("auth_token")
+    num_speakers = args.pop("num_speakers")
+    min_speakers = args.pop("min_speakers")
+    max_speakers = args.pop("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(Diarization(auth_token=auth_token, 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

@@ -138,4 +138,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,
 }

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,188 @@
+import argparse
+import json
+import os
+from pathlib import Path
+import tempfile
+from typing import TYPE_CHECKING, List
+import torch
+
+import ffmpeg
+
+from src.diarization.transcriptLoader import load_transcript
+from src.utils import write_srt
+
+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, **kwargs):
+        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
+        self.pipeline_kwargs = kwargs
+
+    @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/speaker-diarization@2.1", use_auth_token=self.auth_token)
+
+        # 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):
+        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, **self.pipeline_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():
+    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, num_speakers=args.num_speakers, min_speakers=args.min_speakers, max_speakers=args.max_speakers)
+    diarization_result = list(diarization.run(args.audio_file))
+
+    # 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()

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