Ensure GPU memory in diarization can be cleaned up

app.py

@@ -15,6 +15,7 @@ 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
@@ -74,7 +75,10 @@ class WhisperTranscriber:
         self.deleteUploadedFiles = delete_uploaded_files
         self.output_dir = output_dir
 
-        self.diarization: Diarization = None
+        # 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):
@@ -88,6 +92,17 @@ 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, 
@@ -108,9 +123,9 @@ class WhisperTranscriber:
         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)
+            self.set_diarization(auth_token=self.app_config.auth_token, num_speakers=diarization_speakers)
         else:
-            self.diarization = None
+            self.unset_diarization()
 
         return self.transcribe_webui(modelName, languageName, urlData, multipleFiles, microphoneData, task, vadOptions, 
                                      word_timestamps=word_timestamps, highlight_words=highlight_words, progress=progress)
@@ -157,10 +172,10 @@ class WhisperTranscriber:
 
         # 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)
+            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.diarization = None
+            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,
@@ -226,9 +241,9 @@ class WhisperTranscriber:
                     current_progress += source_audio_duration
 
                     # Diarization
-                    if self.diarization:
+                    if self.diarization and self.diarization_kwargs:
                         print("Diarizing ", source.source_path)
-                        diarization_result = list(self.diarization.run(source.source_path))
+                        diarization_result = list(self.diarization.run(source.source_path, **self.diarization_kwargs))
 
                         # Print result
                         print("Diarization result: ")
@@ -494,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, 

cli.py

@@ -162,7 +162,7 @@ def cli():
     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))
+        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)

src/diarization/diarization.py

@@ -1,4 +1,5 @@
 import argparse
+import gc
 import json
 import os
 from pathlib import Path
@@ -8,9 +9,6 @@ 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
@@ -28,7 +26,7 @@ class DiarizationEntry:
         }
 
 class Diarization:
-    def __init__(self, auth_token=None, **kwargs):
+    def __init__(self, auth_token=None):
         if auth_token is None:
             auth_token = os.environ.get("HK_ACCESS_TOKEN")
             if auth_token is None:
@@ -37,7 +35,6 @@ class Diarization:
         self.auth_token = auth_token
         self.initialized = False
         self.pipeline = None
-        self.pipeline_kwargs = kwargs
 
     @staticmethod
     def has_libraries():
@@ -54,6 +51,7 @@ class Diarization:
         from pyannote.audio import Pipeline
 
         self.pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization@2.1", use_auth_token=self.auth_token)
+        self.initialized = True
 
         # Load GPU mode if available
         device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -63,7 +61,7 @@ class Diarization:
         else:
             print("Diarization - using CPU")
 
-    def run(self, audio_file):
+    def run(self, audio_file, **kwargs):
         self.initialize()
         audio_file_obj = Path(audio_file)
 
@@ -78,7 +76,7 @@ class Diarization:
             except ffmpeg.Error as e:
                 print(f"Error occurred during audio conversion: {e.stderr}")
 
-        diarization = self.pipeline(target_file, **self.pipeline_kwargs)
+        diarization = self.pipeline(target_file, **kwargs)
 
         if target_file != audio_file:
             # Delete temp file
@@ -148,6 +146,9 @@ def _write_file(input_file: str, output_path: str, output_extension: str, file_w
     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')
@@ -166,8 +167,8 @@ def main():
     # 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))
+    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:")
@@ -185,4 +186,10 @@ def main():
                 lambda f: write_srt(marked_whisper_result["segments"], f, maxLineWidth=args.max_line_width))
 
 if __name__ == "__main__":
-    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