Spaces:
Sleeping
Sleeping
| 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") |