NeuralInternet/Audio-to-Text_Advanced_Playground
0
1from datetime import datetime2import math3from typing import Iterator4import argparse5 6from io import StringIO7import os8import pathlib9import tempfile10import zipfile11import numpy as np12 13import torch14from src.modelCache import ModelCache15from src.source import get_audio_source_collection16from src.vadParallel import ParallelContext, ParallelTranscription17 18# External programs19import ffmpeg20 21# UI22import gradio as gr23 24from src.download import ExceededMaximumDuration, download_url25from src.utils import slugify, write_srt, write_vtt26from src.vad import AbstractTranscription, NonSpeechStrategy, PeriodicTranscriptionConfig, TranscriptionConfig, VadPeriodicTranscription, VadSileroTranscription27from src.whisperContainer import WhisperContainer28 29# Limitations (set to -1 to disable)30DEFAULT_INPUT_AUDIO_MAX_DURATION = 600 # seconds31 32# Whether or not to automatically delete all uploaded files, to save disk space33DELETE_UPLOADED_FILES = True34 35# Gradio seems to truncate files without keeping the extension, so we need to truncate the file prefix ourself 36MAX_FILE_PREFIX_LENGTH = 1737 38# Limit auto_parallel to a certain number of CPUs (specify vad_cpu_cores to get a higher number)39MAX_AUTO_CPU_CORES = 840 41LANGUAGES = [ 42 "English", "Chinese", "German", "Spanish", "Russian", "Korean", 43 "French", "Japanese", "Portuguese", "Turkish", "Polish", "Catalan", 44 "Dutch", "Arabic", "Swedish", "Italian", "Indonesian", "Hindi", 45 "Finnish", "Vietnamese", "Hebrew", "Ukrainian", "Greek", "Malay", 46 "Czech", "Romanian", "Danish", "Hungarian", "Tamil", "Norwegian", 47 "Thai", "Urdu", "Croatian", "Bulgarian", "Lithuanian", "Latin", 48 "Maori", "Malayalam", "Welsh", "Slovak", "Telugu", "Persian", 49 "Latvian", "Bengali", "Serbian", "Azerbaijani", "Slovenian", 50 "Kannada", "Estonian", "Macedonian", "Breton", "Basque", "Icelandic", 51 "Armenian", "Nepali", "Mongolian", "Bosnian", "Kazakh", "Albanian",52 "Swahili", "Galician", "Marathi", "Punjabi", "Sinhala", "Khmer", 53 "Shona", "Yoruba", "Somali", "Afrikaans", "Occitan", "Georgian", 54 "Belarusian", "Tajik", "Sindhi", "Gujarati", "Amharic", "Yiddish", 55 "Lao", "Uzbek", "Faroese", "Haitian Creole", "Pashto", "Turkmen", 56 "Nynorsk", "Maltese", "Sanskrit", "Luxembourgish", "Myanmar", "Tibetan",57 "Tagalog", "Malagasy", "Assamese", "Tatar", "Hawaiian", "Lingala", 58 "Hausa", "Bashkir", "Javanese", "Sundanese"59]60 61WHISPER_MODELS = ["tiny", "base", "small", "medium", "large", "large-v1", "large-v2"]62 63class WhisperTranscriber:64 def __init__(self, input_audio_max_duration: float = DEFAULT_INPUT_AUDIO_MAX_DURATION, vad_process_timeout: float = None, 65 vad_cpu_cores: int = 1, delete_uploaded_files: bool = DELETE_UPLOADED_FILES, output_dir: str = None):66 self.model_cache = ModelCache()67 self.parallel_device_list = None68 self.gpu_parallel_context = None69 self.cpu_parallel_context = None70 self.vad_process_timeout = vad_process_timeout71 self.vad_cpu_cores = vad_cpu_cores72 73 self.vad_model = None74 self.inputAudioMaxDuration = input_audio_max_duration75 self.deleteUploadedFiles = delete_uploaded_files76 self.output_dir = output_dir77 78 def set_parallel_devices(self, vad_parallel_devices: str):79 self.parallel_device_list = [ device.strip() for device in vad_parallel_devices.split(",") ] if vad_parallel_devices else None80 81 def set_auto_parallel(self, auto_parallel: bool):82 if auto_parallel:83 if torch.cuda.is_available():84 self.parallel_device_list = [ str(gpu_id) for gpu_id in range(torch.cuda.device_count())]85 86 self.vad_cpu_cores = min(os.cpu_count(), MAX_AUTO_CPU_CORES)87 print("[Auto parallel] Using GPU devices " + str(self.parallel_device_list) + " and " + str(self.vad_cpu_cores) + " CPU cores for VAD/transcription.")88 89 # Entry function for the simple tab90 def transcribe_webui_simple(self, modelName, languageName, urlData, multipleFiles, microphoneData, task, vad, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow):91 return self.transcribe_webui(modelName, languageName, urlData, multipleFiles, microphoneData, task, vad, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow)92 93 # Entry function for the full tab94 def transcribe_webui_full(self, modelName, languageName, urlData, multipleFiles, microphoneData, task, vad, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow, 95 initial_prompt: str, temperature: float, best_of: int, beam_size: int, patience: float, length_penalty: float, suppress_tokens: str, 96 condition_on_previous_text: bool, fp16: bool, temperature_increment_on_fallback: float, 97 compression_ratio_threshold: float, logprob_threshold: float, no_speech_threshold: float):98 99 # Handle temperature_increment_on_fallback100 if temperature_increment_on_fallback is not None:101 temperature = tuple(np.arange(temperature, 1.0 + 1e-6, temperature_increment_on_fallback))102 else:103 temperature = [temperature]104 105 return self.transcribe_webui(modelName, languageName, urlData, multipleFiles, microphoneData, task, vad, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow, 106 initial_prompt=initial_prompt, temperature=temperature, best_of=best_of, beam_size=beam_size, patience=patience, length_penalty=length_penalty, suppress_tokens=suppress_tokens,107 condition_on_previous_text=condition_on_previous_text, fp16=fp16,108 compression_ratio_threshold=compression_ratio_threshold, logprob_threshold=logprob_threshold, no_speech_threshold=no_speech_threshold)109 110 def transcribe_webui(self, modelName, languageName, urlData, multipleFiles, microphoneData, task, vad, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow, **decodeOptions: dict):111 try:112 sources = self.__get_source(urlData, multipleFiles, microphoneData)113 114 try:115 selectedLanguage = languageName.lower() if len(languageName) > 0 else None116 selectedModel = modelName if modelName is not None else "base"117 118 model = WhisperContainer(model_name=selectedModel, cache=self.model_cache)119 120 # Result121 download = []122 zip_file_lookup = {}123 text = ""124 vtt = ""125 126 # Write result127 downloadDirectory = tempfile.mkdtemp()128 source_index = 0129 130 outputDirectory = self.output_dir if self.output_dir is not None else downloadDirectory131 132 # Execute whisper133 for source in sources:134 source_prefix = ""135 136 if (len(sources) > 1):137 # Prefix (minimum 2 digits)138 source_index += 1139 source_prefix = str(source_index).zfill(2) + "_"140 print("Transcribing ", source.source_path)141 142 # Transcribe143 result = self.transcribe_file(model, source.source_path, selectedLanguage, task, vad, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow, **decodeOptions)144 filePrefix = slugify(source_prefix + source.get_short_name(), allow_unicode=True)145 146 source_download, source_text, source_vtt = self.write_result(result, filePrefix, outputDirectory)147 148 if len(sources) > 1:149 # Add new line separators150 if (len(source_text) > 0):151 source_text += os.linesep + os.linesep152 if (len(source_vtt) > 0):153 source_vtt += os.linesep + os.linesep154 155 # Append file name to source text too156 source_text = source.get_full_name() + ":" + os.linesep + source_text157 source_vtt = source.get_full_name() + ":" + os.linesep + source_vtt158 159 # Add to result160 download.extend(source_download)161 text += source_text162 vtt += source_vtt163 164 if (len(sources) > 1):165 # Zip files support at least 260 characters, but we'll play it safe and use 200166 zipFilePrefix = slugify(source_prefix + source.get_short_name(max_length=200), allow_unicode=True)167 168 # File names in ZIP file can be longer169 for source_download_file in source_download:170 # Get file postfix (after last -)171 filePostfix = os.path.basename(source_download_file).split("-")[-1]172 zip_file_name = zipFilePrefix + "-" + filePostfix173 zip_file_lookup[source_download_file] = zip_file_name174 175 # Create zip file from all sources176 if len(sources) > 1:177 downloadAllPath = os.path.join(downloadDirectory, "All_Output-" + datetime.now().strftime("%Y%m%d-%H%M%S") + ".zip")178 179 with zipfile.ZipFile(downloadAllPath, 'w', zipfile.ZIP_DEFLATED) as zip:180 for download_file in download:181 # Get file name from lookup182 zip_file_name = zip_file_lookup.get(download_file, os.path.basename(download_file))183 zip.write(download_file, arcname=zip_file_name)184 185 download.insert(0, downloadAllPath)186 187 return download, text, vtt188 189 finally:190 # Cleanup source191 if self.deleteUploadedFiles:192 for source in sources:193 print("Deleting source file " + source.source_path)194 195 try:196 os.remove(source.source_path)197 except Exception as e:198 # Ignore error - it's just a cleanup199 print("Error deleting source file " + source.source_path + ": " + str(e))200 201 except ExceededMaximumDuration as e:202 return [], ("[ERROR]: Maximum remote video length is " + str(e.maxDuration) + "s, file was " + str(e.videoDuration) + "s"), "[ERROR]"203 204 def transcribe_file(self, model: WhisperContainer, audio_path: str, language: str, task: str = None, vad: str = None, 205 vadMergeWindow: float = 5, vadMaxMergeSize: float = 150, vadPadding: float = 1, vadPromptWindow: float = 1, **decodeOptions: dict):206 207 initial_prompt = decodeOptions.pop('initial_prompt', None)208 209 if ('task' in decodeOptions):210 task = decodeOptions.pop('task')211 212 # Callable for processing an audio file213 whisperCallable = model.create_callback(language, task, initial_prompt, **decodeOptions)214 215 # The results216 if (vad == 'silero-vad'):217 # Silero VAD where non-speech gaps are transcribed218 process_gaps = self._create_silero_config(NonSpeechStrategy.CREATE_SEGMENT, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow)219 result = self.process_vad(audio_path, whisperCallable, self.vad_model, process_gaps)220 elif (vad == 'silero-vad-skip-gaps'):221 # Silero VAD where non-speech gaps are simply ignored222 skip_gaps = self._create_silero_config(NonSpeechStrategy.SKIP, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow)223 result = self.process_vad(audio_path, whisperCallable, self.vad_model, skip_gaps)224 elif (vad == 'silero-vad-expand-into-gaps'):225 # Use Silero VAD where speech-segments are expanded into non-speech gaps226 expand_gaps = self._create_silero_config(NonSpeechStrategy.EXPAND_SEGMENT, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow)227 result = self.process_vad(audio_path, whisperCallable, self.vad_model, expand_gaps)228 elif (vad == 'periodic-vad'):229 # Very simple VAD - mark every 5 minutes as speech. This makes it less likely that Whisper enters an infinite loop, but230 # it may create a break in the middle of a sentence, causing some artifacts.231 periodic_vad = VadPeriodicTranscription()232 period_config = PeriodicTranscriptionConfig(periodic_duration=vadMaxMergeSize, max_prompt_window=vadPromptWindow)233 result = self.process_vad(audio_path, whisperCallable, periodic_vad, period_config)234 235 else:236 if (self._has_parallel_devices()):237 # Use a simple period transcription instead, as we need to use the parallel context238 periodic_vad = VadPeriodicTranscription()239 period_config = PeriodicTranscriptionConfig(periodic_duration=math.inf, max_prompt_window=1)240 241 result = self.process_vad(audio_path, whisperCallable, periodic_vad, period_config)242 else:243 # Default VAD244 result = whisperCallable.invoke(audio_path, 0, None, None)245 246 return result247 248 def process_vad(self, audio_path, whisperCallable, vadModel: AbstractTranscription, vadConfig: TranscriptionConfig):249 if (not self._has_parallel_devices()):250 # No parallel devices, so just run the VAD and Whisper in sequence251 return vadModel.transcribe(audio_path, whisperCallable, vadConfig)252 253 gpu_devices = self.parallel_device_list254 255 if (gpu_devices is None or len(gpu_devices) == 0):256 # No GPU devices specified, pass the current environment variable to the first GPU process. This may be NULL.257 gpu_devices = [os.environ.get("CUDA_VISIBLE_DEVICES", None)]258 259 # Create parallel context if needed260 if (self.gpu_parallel_context is None):261 # Create a context wih processes and automatically clear the pool after 1 hour of inactivity262 self.gpu_parallel_context = ParallelContext(num_processes=len(gpu_devices), auto_cleanup_timeout_seconds=self.vad_process_timeout)263 # We also need a CPU context for the VAD264 if (self.cpu_parallel_context is None):265 self.cpu_parallel_context = ParallelContext(num_processes=self.vad_cpu_cores, auto_cleanup_timeout_seconds=self.vad_process_timeout)266 267 parallel_vad = ParallelTranscription()268 return parallel_vad.transcribe_parallel(transcription=vadModel, audio=audio_path, whisperCallable=whisperCallable, 269 config=vadConfig, cpu_device_count=self.vad_cpu_cores, gpu_devices=gpu_devices, 270 cpu_parallel_context=self.cpu_parallel_context, gpu_parallel_context=self.gpu_parallel_context) 271 272 def _has_parallel_devices(self):273 return (self.parallel_device_list is not None and len(self.parallel_device_list) > 0) or self.vad_cpu_cores > 1274 275 def _concat_prompt(self, prompt1, prompt2):276 if (prompt1 is None):277 return prompt2278 elif (prompt2 is None):279 return prompt1280 else:281 return prompt1 + " " + prompt2282 283 def _create_silero_config(self, non_speech_strategy: NonSpeechStrategy, vadMergeWindow: float = 5, vadMaxMergeSize: float = 150, vadPadding: float = 1, vadPromptWindow: float = 1):284 # Use Silero VAD 285 if (self.vad_model is None):286 self.vad_model = VadSileroTranscription()287 288 config = TranscriptionConfig(non_speech_strategy = non_speech_strategy, 289 max_silent_period=vadMergeWindow, max_merge_size=vadMaxMergeSize, 290 segment_padding_left=vadPadding, segment_padding_right=vadPadding, 291 max_prompt_window=vadPromptWindow)292 293 return config294 295 def write_result(self, result: dict, source_name: str, output_dir: str):296 if not os.path.exists(output_dir):297 os.makedirs(output_dir)298 299 text = result["text"]300 language = result["language"]301 languageMaxLineWidth = self.__get_max_line_width(language)302 303 print("Max line width " + str(languageMaxLineWidth))304 vtt = self.__get_subs(result["segments"], "vtt", languageMaxLineWidth)305 srt = self.__get_subs(result["segments"], "srt", languageMaxLineWidth)306 307 output_files = []308 output_files.append(self.__create_file(srt, output_dir, source_name + "-subs.srt"));309 output_files.append(self.__create_file(vtt, output_dir, source_name + "-subs.vtt"));310 output_files.append(self.__create_file(text, output_dir, source_name + "-transcript.txt"));311 312 return output_files, text, vtt313 314 def clear_cache(self):315 self.model_cache.clear()316 self.vad_model = None317 318 def __get_source(self, urlData, multipleFiles, microphoneData):319 return get_audio_source_collection(urlData, multipleFiles, microphoneData, self.inputAudioMaxDuration)320 321 def __get_max_line_width(self, language: str) -> int:322 if (language and language.lower() in ["japanese", "ja", "chinese", "zh"]):323 # Chinese characters and kana are wider, so limit line length to 40 characters324 return 40325 else:326 # TODO: Add more languages327 # 80 latin characters should fit on a 1080p/720p screen328 return 80329 330 def __get_subs(self, segments: Iterator[dict], format: str, maxLineWidth: int) -> str:331 segmentStream = StringIO()332 333 if format == 'vtt':334 write_vtt(segments, file=segmentStream, maxLineWidth=maxLineWidth)335 elif format == 'srt':336 write_srt(segments, file=segmentStream, maxLineWidth=maxLineWidth)337 else:338 raise Exception("Unknown format " + format)339 340 segmentStream.seek(0)341 return segmentStream.read()342 343 def __create_file(self, text: str, directory: str, fileName: str) -> str:344 # Write the text to a file345 with open(os.path.join(directory, fileName), 'w+', encoding="utf-8") as file:346 file.write(text)347 348 return file.name349 350 def close(self):351 print("Closing parallel contexts")352 self.clear_cache()353 354 if (self.gpu_parallel_context is not None):355 self.gpu_parallel_context.close()356 if (self.cpu_parallel_context is not None):357 self.cpu_parallel_context.close()358 359 360def create_ui(input_audio_max_duration, share=False, server_name: str = None, server_port: int = 7860, 361 default_model_name: str = "medium", default_vad: str = None, vad_parallel_devices: str = None, 362 vad_process_timeout: float = None, vad_cpu_cores: int = 1, auto_parallel: bool = False, 363 output_dir: str = None):364 ui = WhisperTranscriber(input_audio_max_duration, vad_process_timeout, vad_cpu_cores, DELETE_UPLOADED_FILES, output_dir)365 366 # Specify a list of devices to use for parallel processing367 ui.set_parallel_devices(vad_parallel_devices)368 ui.set_auto_parallel(auto_parallel)369 370 ui_description = "Whisper is a general-purpose speech recognition model. It is trained on a large dataset of diverse " 371 ui_description += " audio and is also a multi-task model that can perform multilingual speech recognition "372 ui_description += " as well as speech translation and language identification. "373 374 ui_description += "\n\n\n\nFor longer audio files (>10 minutes) not in English, it is recommended that you select Silero VAD (Voice Activity Detector) in the VAD option."375 376 if input_audio_max_duration > 0:377 ui_description += "\n\n" + "Max audio file length: " + str(input_audio_max_duration) + " s"378 379 ui_article = "Read the [documentation here](https://gitlab.com/aadnk/whisper-webui/-/blob/main/docs/options.md)"380 381 simple_inputs = lambda : [382 gr.Dropdown(choices=WHISPER_MODELS, value=default_model_name, label="Model"),383 gr.Dropdown(choices=sorted(LANGUAGES), label="Language"),384 gr.Text(label="URL (YouTube, etc.)"),385 gr.File(label="Upload Files", file_count="multiple"),386 gr.Audio(source="microphone", type="filepath", label="Microphone Input"),387 gr.Dropdown(choices=["transcribe", "translate"], label="Task"),388 gr.Dropdown(choices=["none", "silero-vad", "silero-vad-skip-gaps", "silero-vad-expand-into-gaps", "periodic-vad"], value=default_vad, label="VAD"),389 gr.Number(label="VAD - Merge Window (s)", precision=0, value=5),390 gr.Number(label="VAD - Max Merge Size (s)", precision=0, value=30),391 gr.Number(label="VAD - Padding (s)", precision=None, value=1),392 gr.Number(label="VAD - Prompt Window (s)", precision=None, value=3)393 ]394 395 simple_transcribe = gr.Interface(fn=ui.transcribe_webui_simple, description=ui_description, article=ui_article, inputs=simple_inputs(), outputs=[396 gr.File(label="Download"),397 gr.Text(label="Transcription"), 398 gr.Text(label="Segments")399 ])400 401 full_description = ui_description + "\n\n\n\n" + "Be careful when changing some of the options in the full interface - this can cause the model to crash."402 403 full_transcribe = gr.Interface(fn=ui.transcribe_webui_full, description=full_description, article=ui_article, inputs=[404 *simple_inputs(),405 gr.TextArea(label="Initial Prompt"),406 gr.Number(label="Temperature", value=0),407 gr.Number(label="Best Of - Non-zero temperature", value=5, precision=0),408 gr.Number(label="Beam Size - Zero temperature", value=5, precision=0),409 gr.Number(label="Patience - Zero temperature", value=None),410 gr.Number(label="Length Penalty - Any temperature", value=None), 411 gr.Text(label="Suppress Tokens - Comma-separated list of token IDs", value="-1"),412 gr.Checkbox(label="Condition on previous text", value=True),413 gr.Checkbox(label="FP16", value=True),414 gr.Number(label="Temperature increment on fallback", value=0.2),415 gr.Number(label="Compression ratio threshold", value=2.4),416 gr.Number(label="Logprob threshold", value=-1.0),417 gr.Number(label="No speech threshold", value=0.6)418 ], outputs=[419 gr.File(label="Download"),420 gr.Text(label="Transcription"), 421 gr.Text(label="Segments")422 ])423 424 demo = gr.TabbedInterface([simple_transcribe, full_transcribe], tab_names=["Simple", "Full"])425 426 demo.launch(share=share, server_name=server_name, server_port=server_port)427 428 # Clean up429 ui.close()430 431if __name__ == '__main__':432 parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)433 parser.add_argument("--input_audio_max_duration", type=int, default=DEFAULT_INPUT_AUDIO_MAX_DURATION, help="Maximum audio file length in seconds, or -1 for no limit.")434 parser.add_argument("--share", type=bool, default=False, help="True to share the app on HuggingFace.")435 parser.add_argument("--server_name", type=str, default=None, help="The host or IP to bind to. If None, bind to localhost.")436 parser.add_argument("--server_port", type=int, default=7860, help="The port to bind to.")437 parser.add_argument("--default_model_name", type=str, choices=WHISPER_MODELS, default="medium", help="The default model name.")438 parser.add_argument("--default_vad", type=str, default="silero-vad", help="The default VAD.")439 parser.add_argument("--vad_parallel_devices", type=str, default="", help="A commma delimited list of CUDA devices to use for parallel processing. If None, disable parallel processing.")440 parser.add_argument("--vad_cpu_cores", type=int, default=1, help="The number of CPU cores to use for VAD pre-processing.")441 parser.add_argument("--vad_process_timeout", type=float, default="1800", help="The number of seconds before inactivate processes are terminated. Use 0 to close processes immediately, or None for no timeout.")442 parser.add_argument("--auto_parallel", type=bool, default=False, help="True to use all available GPUs and CPU cores for processing. Use vad_cpu_cores/vad_parallel_devices to specify the number of CPU cores/GPUs to use.")443 parser.add_argument("--output_dir", "-o", type=str, default=None, help="directory to save the outputs")444 445 args = parser.parse_args().__dict__446 create_ui(**args)