Merge pull request #37 from jhj0517/implement_faster-whisper

app.py

@@ -3,6 +3,7 @@ import os
 import argparse
 
 from modules.whisper_Inference import WhisperInference
+from modules.faster_whisper_inference import FasterWhisperInference
 from modules.nllb_inference import NLLBInference
 from ui.htmls import *
 from modules.youtube_manager import get_ytmetas
@@ -12,7 +13,11 @@ class App:
     def __init__(self, args):
         self.args = args
         self.app = gr.Blocks(css=CSS, theme=self.args.theme)
-        self.whisper_inf = WhisperInference()
+        self.whisper_inf = WhisperInference() if self.args.disable_faster_whisper else FasterWhisperInference() 
+        if isinstance(self.whisper_inf, FasterWhisperInference):
+            print("Use Faster Whisper implementation")
+        else:
+            print("Use Open AI Whisper implementation")
         self.nllb_inf = NLLBInference()
 
     @staticmethod
@@ -164,6 +169,7 @@ class App:
 
 # Create the parser for command-line arguments
 parser = argparse.ArgumentParser()
+parser.add_argument('--disable_faster_whisper', type=bool, default=False, nargs='?', const=True, help='Disable the faster_whisper implementation. faster_whipser is implemented by https://github.com/guillaumekln/faster-whisper')
 parser.add_argument('--share', type=bool, default=False, nargs='?', const=True, help='Gradio share value')
 parser.add_argument('--server_name', type=str, default=None, help='Gradio server host')
 parser.add_argument('--server_port', type=int, default=None, help='Gradio server port')

modules/faster_whisper_inference.py

@@ -0,0 +1,353 @@
+import os
+
+import tqdm
+import time
+import numpy as np
+from typing import BinaryIO, Union, Tuple
+from datetime import datetime, timedelta
+
+import faster_whisper
+import whisper
+import torch
+import gradio as gr
+
+from .base_interface import BaseInterface
+from modules.subtitle_manager import get_srt, get_vtt, write_file, safe_filename
+from modules.youtube_manager import get_ytdata, get_ytaudio
+
+
+class FasterWhisperInference(BaseInterface):
+    def __init__(self):
+        super().__init__()
+        self.current_model_size = None
+        self.model = None
+        self.available_models = whisper.available_models()
+        self.available_langs = sorted(list(whisper.tokenizer.LANGUAGES.values()))
+        self.translatable_models = ["large", "large-v1", "large-v2"]
+        self.default_beam_size = 5
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    def transcribe_file(self,
+                        fileobjs: list,
+                        model_size: str,
+                        lang: str,
+                        subformat: str,
+                        istranslate: bool,
+                        add_timestamp: bool,
+                        progress=gr.Progress()
+                        ) -> str:
+        """
+        Write subtitle file from Files
+
+        Parameters
+        ----------
+        fileobjs: list
+            List of files to transcribe from gr.Files()
+        model_size: str
+            Whisper model size from gr.Dropdown()
+        lang: str
+            Source language of the file to transcribe from gr.Dropdown()
+        subformat: str
+            Subtitle format to write from gr.Dropdown(). Supported format: [SRT, WebVTT]
+        istranslate: bool
+            Boolean value from gr.Checkbox() that determines whether to translate to English.
+            It's Whisper's feature to translate speech from another language directly into English end-to-end.
+        add_timestamp: bool
+            Boolean value from gr.Checkbox() that determines whether to add a timestamp at the end of the filename.
+        progress: gr.Progress
+            Indicator to show progress directly in gradio.
+
+        Returns
+        ----------
+        String to return to gr.Textbox()
+        """
+        try:
+            if model_size != self.current_model_size or self.model is None:
+                self.initialize_model(model_size=model_size, progress=progress)
+
+            if lang == "Automatic Detection":
+                lang = None
+
+            files_info = {}
+            for fileobj in fileobjs:
+                transcribed_segments, time_for_task = self.transcribe(
+                    audio=fileobj.name,
+                    lang=lang,
+                    istranslate=istranslate,
+                    progress=progress
+                )
+
+                file_name, file_ext = os.path.splitext(os.path.basename(fileobj.orig_name))
+                file_name = safe_filename(file_name)
+                subtitle = self.generate_and_write_subtitle(
+                    file_name=file_name,
+                    transcribed_segments=transcribed_segments,
+                    add_timestamp=add_timestamp,
+                    subformat=subformat
+                )
+                files_info[file_name] = {"subtitle": subtitle, "time_for_task": time_for_task}
+
+            total_result = ''
+            total_time = 0
+            for file_name, info in files_info.items():
+                total_result += '------------------------------------\n'
+                total_result += f'{file_name}\n\n'
+                total_result += f'{info["subtitle"]}'
+                total_time += info["time_for_task"]
+
+            return f"Done in {self.format_time(total_time)}! Subtitle is in the outputs folder.\n\n{total_result}"
+
+        except Exception as e:
+            print(f"Error transcribing file on line {e}")
+        finally:
+            self.release_cuda_memory()
+            self.remove_input_files([fileobj.name for fileobj in fileobjs])
+
+    def transcribe_youtube(self,
+                           youtubelink: str,
+                           model_size: str,
+                           lang: str,
+                           subformat: str,
+                           istranslate: bool,
+                           add_timestamp: bool,
+                           progress=gr.Progress()
+                           ) -> str:
+        """
+        Write subtitle file from Youtube
+
+        Parameters
+        ----------
+        youtubelink: str
+            Link of Youtube to transcribe from gr.Textbox()
+        model_size: str
+            Whisper model size from gr.Dropdown()
+        lang: str
+            Source language of the file to transcribe from gr.Dropdown()
+        subformat: str
+            Subtitle format to write from gr.Dropdown(). Supported format: [SRT, WebVTT]
+        istranslate: bool
+            Boolean value from gr.Checkbox() that determines whether to translate to English.
+            It's Whisper's feature to translate speech from another language directly into English end-to-end.
+        add_timestamp: bool
+            Boolean value from gr.Checkbox() that determines whether to add a timestamp at the end of the filename.
+        progress: gr.Progress
+            Indicator to show progress directly in gradio.
+
+        Returns
+        ----------
+        String to return to gr.Textbox()
+        """
+        try:
+            if model_size != self.current_model_size or self.model is None:
+                self.initialize_model(model_size=model_size, progress=progress)
+
+            if lang == "Automatic Detection":
+                lang = None
+
+            progress(0, desc="Loading Audio from Youtube..")
+            yt = get_ytdata(youtubelink)
+            audio = get_ytaudio(yt)
+
+            transcribed_segments, time_for_task = self.transcribe(
+                audio=audio,
+                lang=lang,
+                istranslate=istranslate,
+                progress=progress
+            )
+
+            progress(1, desc="Completed!")
+
+            file_name = safe_filename(yt.title)
+            subtitle = self.generate_and_write_subtitle(
+                file_name=file_name,
+                transcribed_segments=transcribed_segments,
+                add_timestamp=add_timestamp,
+                subformat=subformat
+            )
+            return f"Done in {self.format_time(time_for_task)}! Subtitle file is in the outputs folder.\n\n{subtitle}"
+        except Exception as e:
+            return f"Error: {str(e)}"
+        finally:
+            yt = get_ytdata(youtubelink)
+            file_path = get_ytaudio(yt)
+            self.release_cuda_memory()
+            self.remove_input_files([file_path])
+
+    def transcribe_mic(self,
+                       micaudio: str,
+                       model_size: str,
+                       lang: str,
+                       subformat: str,
+                       istranslate: bool,
+                       progress=gr.Progress()
+                       ) -> str:
+        """
+        Write subtitle file from microphone
+
+        Parameters
+        ----------
+        micaudio: str
+            Audio file path from gr.Microphone()
+        model_size: str
+            Whisper model size from gr.Dropdown()
+        lang: str
+            Source language of the file to transcribe from gr.Dropdown()
+        subformat: str
+            Subtitle format to write from gr.Dropdown(). Supported format: [SRT, WebVTT]
+        istranslate: bool
+            Boolean value from gr.Checkbox() that determines whether to translate to English.
+            It's Whisper's feature to translate speech from another language directly into English end-to-end.
+        progress: gr.Progress
+            Indicator to show progress directly in gradio.
+
+        Returns
+        ----------
+        String to return to gr.Textbox()
+        """
+        try:
+            if model_size != self.current_model_size or self.model is None:
+                self.initialize_model(model_size=model_size, progress=progress)
+
+            if lang == "Automatic Detection":
+                lang = None
+
+            progress(0, desc="Loading Audio..")
+
+            transcribed_segments, time_for_task = self.transcribe(
+                audio=micaudio,
+                lang=lang,
+                istranslate=istranslate,
+                progress=progress
+            )
+            progress(1, desc="Completed!")
+
+            subtitle = self.generate_and_write_subtitle(
+                file_name="Mic",
+                transcribed_segments=transcribed_segments,
+                add_timestamp=True,
+                subformat=subformat
+            )
+            return f"Done in {self.format_time(time_for_task)}! Subtitle file is in the outputs folder.\n\n{subtitle}"
+        except Exception as e:
+            return f"Error: {str(e)}"
+        finally:
+            self.release_cuda_memory()
+            self.remove_input_files([micaudio])
+
+    def transcribe(self,
+                   audio: Union[str, BinaryIO, np.ndarray],
+                   lang: str,
+                   istranslate: bool,
+                   progress: gr.Progress
+                   ) -> Tuple[list, float]:
+        """
+        transcribe method for faster-whisper.
+
+        Parameters
+        ----------
+        audio: Union[str, BinaryIO, np.ndarray]
+            Audio path or file binary or Audio numpy array
+        lang: str
+            Source language of the file to transcribe from gr.Dropdown()
+        istranslate: bool
+            Boolean value from gr.Checkbox() that determines whether to translate to English.
+            It's Whisper's feature to translate speech from another language directly into English end-to-end.
+        progress: gr.Progress
+            Indicator to show progress directly in gradio.
+
+        Returns
+        ----------
+        segments_result: list[dict]
+            list of dicts that includes start, end timestamps and transcribed text
+        elapsed_time: float
+            elapsed time for transcription
+        """
+        start_time = time.time()
+        segments, info = self.model.transcribe(
+            audio=audio,
+            language=lang,
+            beam_size=self.default_beam_size,
+            task="translate" if istranslate and self.current_model_size in self.translatable_models else "transcribe"
+        )
+        progress(0, desc="Loading audio..")
+        total_frames = self.get_total_frames(audio=audio, progress=progress)
+
+        segments_result = []
+        for segment in segments:
+            progress(segment.seek / total_frames, desc="Transcribing..")
+            segments_result.append({
+                "start": segment.start,
+                "end": segment.end,
+                "text": segment.text
+            })
+
+        elapsed_time = time.time() - start_time
+        return segments_result, elapsed_time
+
+    def initialize_model(self,
+                         model_size: str,
+                         progress: gr.Progress
+                         ):
+        """
+        Initialize model if it doesn't match with current model size
+        """
+        progress(0, desc="Initializing Model..")
+        self.current_model_size = model_size
+        self.model = faster_whisper.WhisperModel(
+            device=self.device,
+            model_size_or_path=model_size,
+            download_root=os.path.join("models", "Whisper", "faster-whisper"),
+            compute_type="float16"
+        )
+
+    def get_total_frames(self,
+                         audio: Union[str, BinaryIO],
+                         progress: gr.Progress
+                         ) -> float:
+        """
+        This method is only for tracking the progress.
+        Returns total frames to track progress.
+        """
+        progress(0, desc="Loading audio..")
+        decoded_audio = faster_whisper.decode_audio(audio)
+        features = self.model.feature_extractor(decoded_audio)
+        content_frames = features.shape[-1] - self.model.feature_extractor.nb_max_frames
+        return content_frames
+
+    @staticmethod
+    def generate_and_write_subtitle(file_name: str,
+                                    transcribed_segments: list,
+                                    add_timestamp: bool,
+                                    subformat: str,
+                                    ) -> str:
+        """
+        This method writes subtitle file and returns str to gr.Textbox
+        """
+        timestamp = datetime.now().strftime("%m%d%H%M%S")
+        if add_timestamp:
+            output_path = os.path.join("outputs", f"{file_name}-{timestamp}")
+        else:
+            output_path = os.path.join("outputs", f"{file_name}")
+
+        if subformat == "SRT":
+            subtitle = get_srt(transcribed_segments)
+            write_file(subtitle, f"{output_path}.srt")
+        elif subformat == "WebVTT":
+            subtitle = get_vtt(transcribed_segments)
+            write_file(subtitle, f"{output_path}.vtt")
+        return subtitle
+
+    @staticmethod
+    def format_time(elapsed_time: float) -> str:
+        hours, rem = divmod(elapsed_time, 3600)
+        minutes, seconds = divmod(rem, 60)
+
+        time_str = ""
+        if hours:
+            time_str += f"{hours} hours "
+        if minutes:
+            time_str += f"{minutes} minutes "
+        seconds = round(seconds)
+        time_str += f"{seconds} seconds"
+
+        return time_str.strip()

requirements.txt

@@ -1,5 +1,6 @@
 --extra-index-url https://download.pytorch.org/whl/cu117
 torch
 git+https://github.com/jhj0517/jhj0517-whisper.git
+faster-whisper
 gradio==3.37.0
 pytube

user-start-webui.bat

@@ -1,4 +1,5 @@
 :: This batch file is for launching with command line args
+:: See the wiki for a guide to command line arguments: https://github.com/jhj0517/Whisper-WebUI/wiki/Command-Line-Arguments
 @echo off
 
 :: Set values
@@ -8,6 +9,7 @@ set USERNAME=
 set PASSWORD=
 set SHARE=
 set THEME=
+set DISABLE_FASTER_WHISPER=
 
 :: Set args accordingly
 if not "%SERVER_NAME%"=="" (
@@ -28,7 +30,10 @@ if /I "%SHARE%"=="true" (
 if not "%THEME%"=="" (
     set THEME_ARG=--theme %THEME%
 )
+if /I "%DISABLE_FASTER_WHISPER%"=="true" (
+    set DISABLE_FASTER_WHISPER_ARG=--disable_faster_whisper
+)
 
 :: Call the original .bat script with optional arguments
-start-webui.bat %SERVER_NAME_ARG% %SERVER_PORT_ARG% %USERNAME_ARG% %PASSWORD_ARG% %SHARE_ARG% %THEME_ARG%
+start-webui.bat %SERVER_NAME_ARG% %SERVER_PORT_ARG% %USERNAME_ARG% %PASSWORD_ARG% %SHARE_ARG% %THEME_ARG% %DISABLE_FASTER_WHISPER_ARG%
 pause