Update app.py

app.py

@@ -1,88 +1,254 @@
+from typing import Iterator
+
+from io import StringIO
 import os
-import gradio as gr
+import pathlib
+import tempfile
+
+# External programs
 import whisper
-from whisper import tokenizer
-import time
-
-current_size = 'base'
-model = whisper.load_model(current_size)
-AUTO_DETECT_LANG = "Auto Detection"
-
-def transcribe(audio, state={}, model_size='base', delay=1.2, lang=None, translate=False):
-    time.sleep(delay - 1)
-
-    global current_size
-    global model
-    if model_size != current_size:
-        current_size = model_size
-        model = whisper.load_model(current_size)
-  
-    transcription = model.transcribe(
-        audio,
-        language = lang if lang != AUTO_DETECT_LANG else None
-    )
-    state['transcription'] += transcription['text'] + " "
-
-    if translate:
-        x = whisper.load_audio(audio)
-        x = whisper.pad_or_trim(x)
-        mel = whisper.log_mel_spectrogram(x).to(model.device)
-
-        options = whisper.DecodingOptions(task = "translation")
-        translation = whisper.decode(model, mel, options)
-
-        state['translation'] += translation.text + " "
-
-    return state['transcription'], state['translation'], state, f"detected language: {transcription['language']}"
-  
-
-title = "OpenAI's Whisper Real-time Demo"
-description = "A simple demo of OpenAI's [**Whisper**](https://github.com/openai/whisper) speech recognition model. This demo runs on a CPU. For faster inference choose 'tiny' model size and set the language explicitly."
-
-model_size = gr.Dropdown(label="Model size", choices=['base', 'tiny', 'small', 'medium', 'large'], value='base')
-
-delay_slider = gr.inputs.Slider(minimum=1, maximum=5, default=1.2, label="Rate of transcription")
-
-available_languages = sorted(tokenizer.TO_LANGUAGE_CODE.keys())
-available_languages = [lang.capitalize() for lang in available_languages]
-available_languages = [AUTO_DETECT_LANG]+available_languages
-
-lang_dropdown = gr.inputs.Dropdown(choices=available_languages, label="Language", default=AUTO_DETECT_LANG, type="value")
-
-if lang_dropdown==AUTO_DETECT_LANG:
-    lang_dropdown=None
-
-translate_checkbox = gr.inputs.Checkbox(label="Translate to English", default=False)
-
-
-
-transcription_tb = gr.Textbox(label="Transcription", lines=10, max_lines=20)
-translation_tb = gr.Textbox(label="Translation", lines=10, max_lines=20)
-detected_lang = gr.outputs.HTML(label="Detected Language")
-
-state = gr.State({"transcription": "", "translation": ""})
-
-gr.Interface(
-    fn=transcribe,
-    inputs=[
-        gr.Audio(source="microphone", type="filepath", streaming=True),
-        state,
-        model_size,
-        delay_slider,
-        lang_dropdown,
-        translate_checkbox
-        ], 
-    outputs=[
-        transcription_tb,
-        translation_tb,
-        state,
-        detected_lang
-    ],
-    live=True,
-    allow_flagging='never',
-    title=title,
-    description=description,
-).launch(
-    # enable_queue=True,
-    # debug=True
-  )
+import ffmpeg
+
+# UI
+import gradio as gr
+
+from src.download import ExceededMaximumDuration, download_url
+from src.utils import slugify, write_srt, write_vtt
+from src.vad import NonSpeechStrategy, PeriodicTranscriptionConfig, TranscriptionConfig, VadPeriodicTranscription, VadSileroTranscription
+
+# Limitations (set to -1 to disable)
+DEFAULT_INPUT_AUDIO_MAX_DURATION = 600 # seconds
+
+# Whether or not to automatically delete all uploaded files, to save disk space
+DELETE_UPLOADED_FILES = True
+
+# Gradio seems to truncate files without keeping the extension, so we need to truncate the file prefix ourself 
+MAX_FILE_PREFIX_LENGTH = 17
+
+LANGUAGES = [ 
+ "English", "Chinese", "German", "Spanish", "Russian", "Korean", 
+ "French", "Japanese", "Portuguese", "Turkish", "Polish", "Catalan", 
+ "Dutch", "Arabic", "Swedish", "Italian", "Indonesian", "Hindi", 
+ "Finnish", "Vietnamese", "Hebrew", "Ukrainian", "Greek", "Malay", 
+ "Czech", "Romanian", "Danish", "Hungarian", "Tamil", "Norwegian", 
+ "Thai", "Urdu", "Croatian", "Bulgarian", "Lithuanian", "Latin", 
+ "Maori", "Malayalam", "Welsh", "Slovak", "Telugu", "Persian", 
+ "Latvian", "Bengali", "Serbian", "Azerbaijani", "Slovenian", 
+ "Kannada", "Estonian", "Macedonian", "Breton", "Basque", "Icelandic", 
+ "Armenian", "Nepali", "Mongolian", "Bosnian", "Kazakh", "Albanian",
+ "Swahili", "Galician", "Marathi", "Punjabi", "Sinhala", "Khmer", 
+ "Shona", "Yoruba", "Somali", "Afrikaans", "Occitan", "Georgian", 
+ "Belarusian", "Tajik", "Sindhi", "Gujarati", "Amharic", "Yiddish", 
+ "Lao", "Uzbek", "Faroese", "Haitian Creole", "Pashto", "Turkmen", 
+ "Nynorsk", "Maltese", "Sanskrit", "Luxembourgish", "Myanmar", "Tibetan",
+ "Tagalog", "Malagasy", "Assamese", "Tatar", "Hawaiian", "Lingala", 
+ "Hausa", "Bashkir", "Javanese", "Sundanese"
+]
+
+class WhisperTranscriber:
+    def __init__(self, inputAudioMaxDuration: float = DEFAULT_INPUT_AUDIO_MAX_DURATION, deleteUploadedFiles: bool = DELETE_UPLOADED_FILES):
+        self.model_cache = dict()
+
+        self.vad_model = None
+        self.inputAudioMaxDuration = inputAudioMaxDuration
+        self.deleteUploadedFiles = deleteUploadedFiles
+
+    def transcribe_webui(self, modelName, languageName, urlData, uploadFile, microphoneData, task, vad, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow):
+        try:
+            source, sourceName = self.__get_source(urlData, uploadFile, microphoneData)
+            
+            try:
+                selectedLanguage = languageName.lower() if len(languageName) > 0 else None
+                selectedModel = modelName if modelName is not None else "base"
+
+                model = self.model_cache.get(selectedModel, None)
+                
+                if not model:
+                    model = whisper.load_model(selectedModel)
+                    self.model_cache[selectedModel] = model
+
+                # Execute whisper
+                result = self.transcribe_file(model, source, selectedLanguage, task, vad, vadMergeWindow, vadMaxMergeSize, vadPadding, vadPromptWindow)
+
+                # Write result
+                downloadDirectory = tempfile.mkdtemp()
+                
+                filePrefix = slugify(sourceName, allow_unicode=True)
+                download, text, vtt = self.write_result(result, filePrefix, downloadDirectory)
+
+                return download, text, vtt
+
+            finally:
+                # Cleanup source
+                if self.deleteUploadedFiles:
+                    print("Deleting source file " + source)
+                    os.remove(source)
+        
+        except ExceededMaximumDuration as e:
+            return [], ("[ERROR]: Maximum remote video length is " + str(e.maxDuration) + "s, file was " + str(e.videoDuration) + "s"), "[ERROR]"
+
+    def transcribe_file(self, model: whisper.Whisper, 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):
+        
+        initial_prompt = decodeOptions.pop('initial_prompt', None)
+
+        if ('task' in decodeOptions):
+            task = decodeOptions.pop('task')
+
+        # Callable for processing an audio file
+        whisperCallable = lambda audio, segment_index, prompt, detected_language : model.transcribe(audio, \
+                 language=language if language else detected_language, task=task, \
+                 initial_prompt=self._concat_prompt(initial_prompt, prompt) if segment_index == 0 else prompt, \
+                 **decodeOptions)
+
+        # 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.vad_model.transcribe(audio_path, whisperCallable, process_gaps)
+        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.vad_model.transcribe(audio_path, whisperCallable, skip_gaps)
+        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.vad_model.transcribe(audio_path, whisperCallable, expand_gaps)
+        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()
+            result = periodic_vad.transcribe(audio_path, whisperCallable, PeriodicTranscriptionConfig(periodic_duration=vadMaxMergeSize, max_prompt_window=vadPromptWindow))
+        else:
+            # Default VAD
+            result = whisperCallable(audio_path, 0, None, None)
+
+        return result
+
+    def _concat_prompt(self, prompt1, prompt2):
+        if (prompt1 is None):
+            return prompt2
+        elif (prompt2 is None):
+            return prompt1
+        else:
+            return prompt1 + " " + prompt2
+
+    def _create_silero_config(self, non_speech_strategy: NonSpeechStrategy, vadMergeWindow: float = 5, vadMaxMergeSize: float = 150, vadPadding: float = 1, vadPromptWindow: float = 1):
+        # Use Silero VAD 
+        if (self.vad_model is None):
+            self.vad_model = VadSileroTranscription()
+
+        config = TranscriptionConfig(non_speech_strategy = non_speech_strategy, 
+                max_silent_period=vadMergeWindow, max_merge_size=vadMaxMergeSize, 
+                segment_padding_left=vadPadding, segment_padding_right=vadPadding, 
+                max_prompt_window=vadPromptWindow)
+
+        return config
+
+    def write_result(self, result: dict, source_name: str, output_dir: str):
+        if not os.path.exists(output_dir):
+            os.makedirs(output_dir)
+
+        text = result["text"]
+        language = result["language"]
+        languageMaxLineWidth = self.__get_max_line_width(language)
+
+        print("Max line width " + str(languageMaxLineWidth))
+        vtt = self.__get_subs(result["segments"], "vtt", languageMaxLineWidth)
+        srt = self.__get_subs(result["segments"], "srt", languageMaxLineWidth)
+
+        output_files = []
+        output_files.append(self.__create_file(srt, output_dir, source_name + "-subs.srt"));
+        output_files.append(self.__create_file(vtt, output_dir, source_name + "-subs.vtt"));
+        output_files.append(self.__create_file(text, output_dir, source_name + "-transcript.txt"));
+
+        return output_files, text, vtt
+
+    def clear_cache(self):
+        self.model_cache = dict()
+        self.vad_model = None
+
+    def __get_source(self, urlData, uploadFile, microphoneData):
+        if urlData:
+            # Download from YouTube
+            source = download_url(urlData, self.inputAudioMaxDuration)[0]
+        else:
+            # File input
+            source = uploadFile if uploadFile is not None else microphoneData
+
+            if self.inputAudioMaxDuration > 0:
+                # Calculate audio length
+                audioDuration = ffmpeg.probe(source)["format"]["duration"]
+            
+                if float(audioDuration) > self.inputAudioMaxDuration:
+                    raise ExceededMaximumDuration(videoDuration=audioDuration, maxDuration=self.inputAudioMaxDuration, message="Video is too long")
+
+        file_path = pathlib.Path(source)
+        sourceName = file_path.stem[:MAX_FILE_PREFIX_LENGTH] + file_path.suffix
+
+        return source, sourceName
+
+    def __get_max_line_width(self, language: str) -> int:
+        if (language and language.lower() in ["japanese", "ja", "chinese", "zh"]):
+            # Chinese characters and kana are wider, so limit line length to 40 characters
+            return 40
+        else:
+            # TODO: Add more languages
+            # 80 latin characters should fit on a 1080p/720p screen
+            return 80
+
+    def __get_subs(self, segments: Iterator[dict], format: str, maxLineWidth: int) -> str:
+        segmentStream = StringIO()
+
+        if format == 'vtt':
+            write_vtt(segments, file=segmentStream, maxLineWidth=maxLineWidth)
+        elif format == 'srt':
+            write_srt(segments, file=segmentStream, maxLineWidth=maxLineWidth)
+        else:
+            raise Exception("Unknown format " + format)
+
+        segmentStream.seek(0)
+        return segmentStream.read()
+
+    def __create_file(self, text: str, directory: str, fileName: str) -> str:
+        # Write the text to a file
+        with open(os.path.join(directory, fileName), 'w+', encoding="utf-8") as file:
+            file.write(text)
+
+        return file.name
+
+
+def create_ui(inputAudioMaxDuration, share=False, server_name: str = None):
+    ui = WhisperTranscriber(inputAudioMaxDuration)
+
+    ui_description = "Whisper是一个语音转文字模型，经过多个语音数据集的训练而成。也可以进行多语言的识别任务和翻译（多种语言翻译成英文）" 
+
+
+    ui_description += "\n\n\n\n对于时长大于10分钟的非英语音频文件，建议选择VAD选项中的Silero VAD (语音活动检测器)。"
+
+    if inputAudioMaxDuration > 0:
+        ui_description += "\n\n" + "Max audio file length: " + str(inputAudioMaxDuration) + " s"
+
+    
+    demo = gr.Interface(fn=ui.transcribe_webui, description=ui_description, inputs=[
+        gr.Dropdown(choices=["tiny", "base", "small", "medium", "large"], value="medium", label="Model"),
+        gr.Dropdown(choices=sorted(LANGUAGES), label="Language"),
+        gr.Text(label="URL (YouTube, etc.)"),
+        gr.Audio(source="upload", type="filepath", label="Upload Audio"), 
+        gr.Audio(source="microphone", type="filepath", label="Microphone Input"),
+        gr.Dropdown(choices=["transcribe", "translate"], label="Task"),
+        gr.Dropdown(choices=["none", "silero-vad", "silero-vad-skip-gaps", "silero-vad-expand-into-gaps", "periodic-vad"], label="VAD"),
+        gr.Number(label="VAD - Merge Window (s)", precision=0, value=5),
+        gr.Number(label="VAD - Max Merge Size (s)", precision=0, value=30),
+        gr.Number(label="VAD - Padding (s)", precision=None, value=1),
+        gr.Number(label="VAD - Prompt Window (s)", precision=None, value=3)
+    ], outputs=[
+        gr.File(label="Download"),
+        gr.Text(label="Transcription"), 
+        gr.Text(label="Segments")
+    ])
+
+    demo.launch(share=share, server_name=server_name)   
+
+if __name__ == '__main__':
+    create_ui(DEFAULT_INPUT_AUDIO_MAX_DURATION)