Make a separate process timeout for diarization

app.py

@@ -33,7 +33,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
@@ -95,7 +95,8 @@ class WhisperTranscriber:
     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)
+                                                    auto_cleanup_timeout_seconds=self.app_config.diarization_process_timeout, 
+                                                    cache=self.model_cache)
         # Set parameters
         self.diarization_kwargs = kwargs
 
@@ -688,6 +689,15 @@ if __name__ == '__main__':
                         help="the compute type to use for inference")
     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")
+    
+    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__
 

config.json5

@@ -148,5 +148,7 @@
     // The minimum number of speakers to detect
     "diarization_min_speakers": 1,
     // The maximum number of speakers to detect
-    "diarization_max_speakers": 5,
+    "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,
 }

src/config.py

@@ -72,7 +72,8 @@ class ApplicationConfig:
                  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_min_speakers: int = 1, diarization_max_speakers: int = 5,
+                 diarization_process_timeout: int = 60):
         
         self.models = models
         
@@ -130,6 +131,7 @@ class ApplicationConfig:
         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/diarizationContainer.py

@@ -16,7 +16,8 @@ class DiarizationContainer:
         # 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)
+            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: