Add progress listener to none/VAD

Note that we don't handle progress of parallel transcription yet.

app.py

@@ -1,6 +1,6 @@
 from datetime import datetime
 import math
-from typing import Iterator
+from typing import Iterator, Union
 import argparse
 
 from io import StringIO
@@ -12,6 +12,7 @@ import numpy as np
 
 import torch
 from src.config import ApplicationConfig
+from src.hooks.whisperProgressHook import ProgressListener, create_progress_listener_handle
 from src.modelCache import ModelCache
 from src.source import get_audio_source_collection
 from src.vadParallel import ParallelContext, ParallelTranscription
@@ -87,14 +88,17 @@ class WhisperTranscriber:
             print("[Auto parallel] Using GPU devices " + str(self.parallel_device_list) + " and " + str(self.vad_cpu_cores) + " CPU cores for VAD/transcription.")
 
     # Entry function for the simple tab
-    def transcribe_webui_simple(self, modelName, languageName, urlData, multipleFiles, microphoneData, task, vad, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow):
-        return self.transcribe_webui(modelName, languageName, urlData, multipleFiles, microphoneData, task, vad, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow)
+    def transcribe_webui_simple(self, modelName, languageName, urlData, multipleFiles, microphoneData, task, vad, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow, 
+                                progress=gr.Progress()):
+        return self.transcribe_webui(modelName, languageName, urlData, multipleFiles, microphoneData, task, vad, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow, 
+                                     progress=progress)
 
     # Entry function for the full tab
     def transcribe_webui_full(self, modelName, languageName, urlData, multipleFiles, microphoneData, task, vad, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow, 
                                     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, 
+                                    progress=gr.Progress()):
 
         # Handle temperature_increment_on_fallback
         if temperature_increment_on_fallback is not None:
@@ -105,9 +109,11 @@ class WhisperTranscriber:
         return self.transcribe_webui(modelName, languageName, urlData, multipleFiles, microphoneData, task, vad, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow, 
                                      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,
-                                     compression_ratio_threshold=compression_ratio_threshold, logprob_threshold=logprob_threshold, no_speech_threshold=no_speech_threshold)
+                                     compression_ratio_threshold=compression_ratio_threshold, logprob_threshold=logprob_threshold, no_speech_threshold=no_speech_threshold, 
+                                     progress=progress)
 
-    def transcribe_webui(self, modelName, languageName, urlData, multipleFiles, microphoneData, task, vad, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow, **decodeOptions: dict):
+    def transcribe_webui(self, modelName, languageName, urlData, multipleFiles, microphoneData, task, vad, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow, 
+                         progress: gr.Progress = None, **decodeOptions: dict):
         try:
             sources = self.__get_source(urlData, multipleFiles, microphoneData)
             
@@ -140,7 +146,7 @@ class WhisperTranscriber:
                         print("Transcribing ", source.source_path)
 
                     # Transcribe
-                    result = self.transcribe_file(model, source.source_path, selectedLanguage, task, vad, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow, **decodeOptions)
+                    result = self.transcribe_file(model, source.source_path, selectedLanguage, task, vad, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow, progress, **decodeOptions)
                     filePrefix = slugify(source_prefix + source.get_short_name(), allow_unicode=True)
 
                     source_download, source_text, source_vtt = self.write_result(result, filePrefix, outputDirectory)
@@ -202,7 +208,8 @@ class WhisperTranscriber:
             return [], ("[ERROR]: Maximum remote video length is " + str(e.maxDuration) + "s, file was " + str(e.videoDuration) + "s"), "[ERROR]"
 
     def transcribe_file(self, model: WhisperContainer, audio_path: str, language: str, task: str = None, vad: str = None, 
-                        vadMergeWindow: float = 5, vadMaxMergeSize: float = 150, vadPadding: float = 1, vadPromptWindow: float = 1, **decodeOptions: dict):
+                        vadMergeWindow: float = 5, vadMaxMergeSize: float = 150, vadPadding: float = 1, vadPromptWindow: float = 1, 
+                        progress: gr.Progress = None, **decodeOptions: dict):
         
         initial_prompt = decodeOptions.pop('initial_prompt', None)
 
@@ -212,25 +219,28 @@ class WhisperTranscriber:
         # Callable for processing an audio file
         whisperCallable = model.create_callback(language, task, initial_prompt, **decodeOptions)
 
+        # A listener that will report progress to Gradio
+        progressListener = self._create_progress_listener(progress)
+
         # The results
         if (vad == 'silero-vad'):
             # Silero VAD where non-speech gaps are transcribed
             process_gaps = self._create_silero_config(NonSpeechStrategy.CREATE_SEGMENT, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow)
-            result = self.process_vad(audio_path, whisperCallable, self.vad_model, process_gaps)
+            result = self.process_vad(audio_path, whisperCallable, self.vad_model, process_gaps, progressListener=progressListener)
         elif (vad == 'silero-vad-skip-gaps'):
             # Silero VAD where non-speech gaps are simply ignored
             skip_gaps = self._create_silero_config(NonSpeechStrategy.SKIP, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow)
-            result = self.process_vad(audio_path, whisperCallable, self.vad_model, skip_gaps)
+            result = self.process_vad(audio_path, whisperCallable, self.vad_model, skip_gaps, progressListener=progressListener)
         elif (vad == 'silero-vad-expand-into-gaps'):
             # Use Silero VAD where speech-segments are expanded into non-speech gaps
             expand_gaps = self._create_silero_config(NonSpeechStrategy.EXPAND_SEGMENT, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow)
-            result = self.process_vad(audio_path, whisperCallable, self.vad_model, expand_gaps)
+            result = self.process_vad(audio_path, whisperCallable, self.vad_model, expand_gaps, progressListener=progressListener)
         elif (vad == 'periodic-vad'):
             # Very simple VAD - mark every 5 minutes as speech. This makes it less likely that Whisper enters an infinite loop, but
             # it may create a break in the middle of a sentence, causing some artifacts.
             periodic_vad = VadPeriodicTranscription()
             period_config = PeriodicTranscriptionConfig(periodic_duration=vadMaxMergeSize, max_prompt_window=vadPromptWindow)
-            result = self.process_vad(audio_path, whisperCallable, periodic_vad, period_config)
+            result = self.process_vad(audio_path, whisperCallable, periodic_vad, period_config, progressListener=progressListener)
 
         else:
             if (self._has_parallel_devices()):
@@ -238,18 +248,38 @@ class WhisperTranscriber:
                 periodic_vad = VadPeriodicTranscription()
                 period_config = PeriodicTranscriptionConfig(periodic_duration=math.inf, max_prompt_window=1)
 
-                result = self.process_vad(audio_path, whisperCallable, periodic_vad, period_config)
+                result = self.process_vad(audio_path, whisperCallable, periodic_vad, period_config, progressListener=progressListener)
             else:
                 # Default VAD
-                result = whisperCallable.invoke(audio_path, 0, None, None)
+                result = whisperCallable.invoke(audio_path, 0, None, None, progress_listener=progressListener)
 
         return result
 
-    def process_vad(self, audio_path, whisperCallable, vadModel: AbstractTranscription, vadConfig: TranscriptionConfig):
+    def _create_progress_listener(self, progress: gr.Progress):
+        if (progress is None):
+            # Dummy progress listener
+            return ProgressListener()
+        
+        class ForwardingProgressListener(ProgressListener):
+            def __init__(self, progress: gr.Progress):
+                self.progress = progress
+
+            def on_progress(self, current: Union[int, float], total: Union[int, float]):
+                # From 0 to 1
+                self.progress(current / total)
+
+            def on_finished(self):
+                self.progress(1)
+
+        return ForwardingProgressListener(progress)
+
+    def process_vad(self, audio_path, whisperCallable, vadModel: AbstractTranscription, vadConfig: TranscriptionConfig, 
+                    progressListener: ProgressListener = None):
         if (not self._has_parallel_devices()):
             # No parallel devices, so just run the VAD and Whisper in sequence
-            return vadModel.transcribe(audio_path, whisperCallable, vadConfig)
+            return vadModel.transcribe(audio_path, whisperCallable, vadConfig, progressListener=progressListener)
 
+        # TODO: Handle progress listener
         gpu_devices = self.parallel_device_list
 
         if (gpu_devices is None or len(gpu_devices) == 0):

src/hooks/whisperProgressHook.py

@@ -0,0 +1,119 @@
+import sys
+import threading
+from typing import List, Union
+import tqdm
+
+class ProgressListener:
+    def on_progress(self, current: Union[int, float], total: Union[int, float]):
+        self.total = total
+
+    def on_finished(self):
+        pass
+
+class ProgressListenerHandle:
+    def __init__(self, listener: ProgressListener):
+        self.listener = listener
+    
+    def __enter__(self):
+        register_thread_local_progress_listener(self.listener)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        unregister_thread_local_progress_listener(self.listener)
+        
+        if exc_type is None:
+            self.listener.on_finished()
+
+class SubTaskProgressListener(ProgressListener):
+    """
+    A sub task listener that reports the progress of a sub task to a base task listener
+    
+    Parameters
+    ----------
+    base_task_listener : ProgressListener
+        The base progress listener to accumulate overall progress in.
+    base_task_total : float
+        The maximum total progress that will be reported to the base progress listener.
+    sub_task_start : float
+        The starting progress of a sub task, in respect to the base progress listener.
+    sub_task_total : float
+        The total amount of progress a sub task will report to the base progress listener.
+    """
+    def __init__(
+        self,
+        base_task_listener: ProgressListener,
+        base_task_total: float,
+        sub_task_start: float,
+        sub_task_total: float,
+    ):
+        self.base_task_listener = base_task_listener
+        self.base_task_total = base_task_total
+        self.sub_task_start = sub_task_start
+        self.sub_task_total = sub_task_total
+
+    def on_progress(self, current: Union[int, float], total: Union[int, float]):
+        sub_task_progress_frac = current / total
+        sub_task_progress = self.sub_task_start + self.sub_task_total * sub_task_progress_frac
+        self.base_task_listener.on_progress(sub_task_progress, self.base_task_total)
+
+    def on_finished(self):
+        self.base_task_listener.on_progress(self.sub_task_start + self.sub_task_total, self.base_task_total)
+
+class _CustomProgressBar(tqdm.tqdm):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._current = self.n  # Set the initial value
+
+    def update(self, n):
+        super().update(n)
+        # Because the progress bar might be disabled, we need to manually update the progress
+        self._current += n
+
+        # Inform listeners
+        listeners = _get_thread_local_listeners()
+
+        for listener in listeners:
+            listener.on_progress(self._current, self.total)
+
+_thread_local = threading.local()
+
+def _get_thread_local_listeners():
+    if not hasattr(_thread_local, 'listeners'):
+        _thread_local.listeners = []
+    return _thread_local.listeners
+
+_hooked = False
+
+def init_progress_hook():
+    global _hooked
+
+    if _hooked:
+        return
+
+    # Inject into tqdm.tqdm of Whisper, so we can see progress
+    import whisper.transcribe 
+    transcribe_module = sys.modules['whisper.transcribe']
+    transcribe_module.tqdm.tqdm = _CustomProgressBar
+    _hooked = True
+
+def register_thread_local_progress_listener(progress_listener: ProgressListener):
+    # This is a workaround for the fact that the progress bar is not exposed in the API
+    init_progress_hook()
+
+    listeners = _get_thread_local_listeners()
+    listeners.append(progress_listener)
+
+def unregister_thread_local_progress_listener(progress_listener: ProgressListener):
+    listeners = _get_thread_local_listeners()
+    
+    if progress_listener in listeners:
+        listeners.remove(progress_listener)
+
+def create_progress_listener_handle(progress_listener: ProgressListener):
+    return ProgressListenerHandle(progress_listener)
+
+if __name__ == '__main__':
+    with create_progress_listener_handle(ProgressListener()) as listener:
+        # Call model.transcribe here
+        pass
+
+    print("Done")

src/vad.py

@@ -5,6 +5,7 @@ import time
 from typing import Any, Deque, Iterator, List, Dict
 
 from pprint import pprint
+from src.hooks.whisperProgressHook import ProgressListener, SubTaskProgressListener, create_progress_listener_handle
 from src.modelCache import GLOBAL_MODEL_CACHE, ModelCache
 
 from src.segments import merge_timestamps
@@ -135,7 +136,8 @@ class AbstractTranscription(ABC):
             pprint(merged)
         return merged
 
-    def transcribe(self, audio: str, whisperCallable: WhisperCallback, config: TranscriptionConfig):
+    def transcribe(self, audio: str, whisperCallable: WhisperCallback, config: TranscriptionConfig, 
+                   progressListener: ProgressListener = None):
         """
         Transcribe the given audo file.
 
@@ -184,7 +186,7 @@ class AbstractTranscription(ABC):
             segment_duration = segment_end - segment_start
 
             if segment_duration < MIN_SEGMENT_DURATION:
-                continue;
+                continue
 
             # Audio to run on Whisper
             segment_audio = self.get_audio_segment(audio, start_time = str(segment_start), duration = str(segment_duration))
@@ -196,7 +198,9 @@ class AbstractTranscription(ABC):
 
             print("Running whisper from ", format_timestamp(segment_start), " to ", format_timestamp(segment_end), ", duration: ", 
                   segment_duration, "expanded: ", segment_expand_amount, "prompt: ", segment_prompt, "language: ", detected_language)
-            segment_result = whisperCallable.invoke(segment_audio, segment_index, segment_prompt, detected_language)
+            
+            scaled_progress_listener = SubTaskProgressListener(progressListener, base_task_total=max_audio_duration, sub_task_start=segment_start, sub_task_total=segment_duration) 
+            segment_result = whisperCallable.invoke(segment_audio, segment_index, segment_prompt, detected_language, progress_listener=scaled_progress_listener)
 
             adjusted_segments = self.adjust_timestamp(segment_result["segments"], adjust_seconds=segment_start, max_source_time=segment_duration)
 
@@ -226,7 +230,7 @@ class AbstractTranscription(ABC):
             result['language'] = detected_language
 
         return result
-            
+    
     def __update_prompt_window(self, prompt_window: Deque, adjusted_segments: List, segment_end: float, segment_gap: bool, config: TranscriptionConfig):
         if (config.max_prompt_window is not None and config.max_prompt_window > 0):
             # Add segments to the current prompt window (unless it is a speech gap)

src/whisperContainer.py

@@ -1,8 +1,13 @@
 # External programs
 import os
+import sys
 from typing import List
+
 import whisper
+from whisper import Whisper
+
 from src.config import ModelConfig
+from src.hooks.whisperProgressHook import ProgressListener, create_progress_listener_handle
 
 from src.modelCache import GLOBAL_MODEL_CACHE, ModelCache
 
@@ -116,7 +121,7 @@ class WhisperCallback:
         self.initial_prompt = initial_prompt
         self.decodeOptions = decodeOptions
         
-    def invoke(self, audio, segment_index: int, prompt: str, detected_language: str):
+    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.
 
@@ -139,10 +144,18 @@ class WhisperCallback:
         """
         model = self.model_container.get_model()
 
+        if progress_listener is not None:
+            with create_progress_listener_handle(progress_listener):
+                return self._transcribe(model, audio, segment_index, prompt, detected_language)
+        else:
+            return self._transcribe(model, audio, segment_index, prompt, detected_language)
+    
+    def _transcribe(self, model: Whisper, audio, segment_index: int, prompt: str, detected_language: str):
         return model.transcribe(audio, \
-                 language=self.language if self.language else detected_language, task=self.task, \
-                 initial_prompt=self._concat_prompt(self.initial_prompt, prompt) if segment_index == 0 else prompt, \
-                 **self.decodeOptions)
+            language=self.language if self.language else detected_language, task=self.task, \
+            initial_prompt=self._concat_prompt(self.initial_prompt, prompt) if segment_index == 0 else prompt, \
+            **self.decodeOptions
+        )
 
     def _concat_prompt(self, prompt1, prompt2):
         if (prompt1 is None):