refactored for better read

app.py

@@ -1,139 +1,158 @@
 import gradio as gr
+import os
+import argparse
+
 from modules.whisper_Inference import WhisperInference
 from modules.nllb_inference import NLLBInference
-import os
 from ui.htmls import *
 from modules.youtube_manager import get_ytmetas
-import argparse
-
-# Create the parser
-parser = argparse.ArgumentParser()
-parser.add_argument('--share', type=bool, default=False, nargs='?', const=True,
-                    help='Share value')
-args = parser.parse_args()
-
-
-def open_folder(folder_path):
-    if os.path.exists(folder_path):
-        os.system(f"start {folder_path}")
-    else:
-        print(f"The folder {folder_path} does not exist.")
-
-
-def on_change_models(model_size):
-    translatable_model = ["large", "large-v1", "large-v2"]
-    if model_size not in translatable_model:
-        return gr.Checkbox.update(visible=False, value=False, interactive=False)
-    else:
-        return gr.Checkbox.update(visible=True, value=False, label="Translate to English?", interactive=True)
 
-whisper_inf = WhisperInference()
-nllb_inf = NLLBInference()
-block = gr.Blocks(css=CSS).queue(api_open=False)
-
-with block:
-    with gr.Row():
-        with gr.Column():
-            gr.Markdown(MARKDOWN, elem_id="md_project")
-    with gr.Tabs():
-        with gr.TabItem("File"):  # tab1
-            with gr.Row():
-                input_file = gr.Files(type="file", label="Upload File here")
-            with gr.Row():
-                dd_model = gr.Dropdown(choices=whisper_inf.available_models, value="large-v2", label="Model")
-                dd_lang = gr.Dropdown(choices=["Automatic Detection"] + whisper_inf.available_langs,
-                                      value="Automatic Detection", label="Language")
-                dd_subformat = gr.Dropdown(["SRT", "WebVTT"], value="SRT", label="Subtitle Format")
-            with gr.Row():
-                cb_translate = gr.Checkbox(value=False, label="Translate to English?", interactive=True)
-            with gr.Row():
-                btn_run = gr.Button("GENERATE SUBTITLE FILE", variant="primary")
-            with gr.Row():
-                tb_indicator = gr.Textbox(label="Output", scale=8)
-                btn_openfolder = gr.Button('📂', scale=2)
 
-            btn_run.click(fn=whisper_inf.transcribe_file,
-                          inputs=[input_file, dd_model, dd_lang, dd_subformat, cb_translate], outputs=[tb_indicator])
-            btn_openfolder.click(fn=lambda: open_folder("outputs"), inputs=None, outputs=None)
-            dd_model.change(fn=on_change_models, inputs=[dd_model], outputs=[cb_translate])
+class App:
+    def __init__(self, args):
+        self.args = args
+        self.app = gr.Blocks(css=CSS)
+        self.whisper_inf = WhisperInference()
+        self.nllb_inf = NLLBInference()
 
-        with gr.TabItem("Youtube"):  # tab2
-            with gr.Row():
-                tb_youtubelink = gr.Textbox(label="Youtube Link")
-            with gr.Row(equal_height=True):
-                with gr.Column():
-                    img_thumbnail = gr.Image(label="Youtube Thumbnail")
-                with gr.Column():
-                    tb_title = gr.Label(label="Youtube Title")
-                    tb_description = gr.Textbox(label="Youtube Description", max_lines=15)
-            with gr.Row():
-                dd_model = gr.Dropdown(choices=whisper_inf.available_models, value="large-v2", label="Model")
-                dd_lang = gr.Dropdown(choices=["Automatic Detection"] + whisper_inf.available_langs,
-                                      value="Automatic Detection", label="Language")
-                dd_subformat = gr.Dropdown(choices=["SRT", "WebVTT"], value="SRT", label="Subtitle Format")
-            with gr.Row():
-                cb_translate = gr.Checkbox(value=False, label="Translate to English?", interactive=True)
-            with gr.Row():
-                btn_run = gr.Button("GENERATE SUBTITLE FILE", variant="primary")
-            with gr.Row():
-                tb_indicator = gr.Textbox(label="Output", scale=8)
-                btn_openfolder = gr.Button('📂', scale=2)
-
-            btn_run.click(fn=whisper_inf.transcribe_youtube,
-                          inputs=[tb_youtubelink, dd_model, dd_lang, dd_subformat, cb_translate],
-                          outputs=[tb_indicator])
-            tb_youtubelink.change(get_ytmetas, inputs=[tb_youtubelink],
-                                  outputs=[img_thumbnail, tb_title, tb_description])
-            btn_openfolder.click(fn=lambda: open_folder("outputs"), inputs=None, outputs=None)
-            dd_model.change(fn=on_change_models, inputs=[dd_model], outputs=[cb_translate])
-
-        with gr.TabItem("Mic"):  # tab3
-            with gr.Row():
-                mic_input = gr.Microphone(label="Record with Mic", type="filepath", interactive=True)
-            with gr.Row():
-                dd_model = gr.Dropdown(choices=whisper_inf.available_models, value="large-v2", label="Model")
-                dd_lang = gr.Dropdown(choices=["Automatic Detection"] + whisper_inf.available_langs,
-                                      value="Automatic Detection", label="Language")
-                dd_subformat = gr.Dropdown(["SRT", "WebVTT"], value="SRT", label="Subtitle Format")
-            with gr.Row():
-                cb_translate = gr.Checkbox(value=False, label="Translate to English?", interactive=True)
-            with gr.Row():
-                btn_run = gr.Button("GENERATE SUBTITLE FILE", variant="primary")
-            with gr.Row():
-                tb_indicator = gr.Textbox(label="Output", scale=8)
-                btn_openfolder = gr.Button('📂', scale=2)
+    @staticmethod
+    def open_folder(folder_path: str):
+        if os.path.exists(folder_path):
+            os.system(f"start {folder_path}")
+        else:
+            print(f"The folder {folder_path} does not exist.")
 
-            btn_run.click(fn=whisper_inf.transcribe_mic,
-                          inputs=[mic_input, dd_model, dd_lang, dd_subformat, cb_translate], outputs=[tb_indicator])
-            btn_openfolder.click(fn=lambda: open_folder("outputs"), inputs=None, outputs=None)
-            dd_model.change(fn=on_change_models, inputs=[dd_model], outputs=[cb_translate])
+    @staticmethod
+    def on_change_models(model_size: str):
+        translatable_model = ["large", "large-v1", "large-v2"]
+        if model_size not in translatable_model:
+            return gr.Checkbox.update(visible=False, value=False, interactive=False)
+        else:
+            return gr.Checkbox.update(visible=True, value=False, label="Translate to English?", interactive=True)
 
-        with gr.TabItem("T2T Translation"):  # tab 4
+    def launch(self):
+        with self.app:
             with gr.Row():
-                file_subs = gr.Files(type="file", label="Upload Subtitle Files to translate here",
-                                     file_types=['.vtt', '.srt'])
-
-            with gr.TabItem("NLLB"):  # sub tab1
-                with gr.Row():
-                    dd_nllb_model = gr.Dropdown(label="Model", value=nllb_inf.default_model_size,
-                                                choices=nllb_inf.available_models)
-                    dd_nllb_sourcelang = gr.Dropdown(label="Source Language", choices=nllb_inf.available_source_langs)
-                    dd_nllb_targetlang = gr.Dropdown(label="Target Language", choices=nllb_inf.available_target_langs)
-                with gr.Row():
-                    btn_run = gr.Button("TRANSLATE SUBTITLE FILE", variant="primary")
-                with gr.Row():
-                    tb_indicator = gr.Textbox(label="Output", scale=8)
-                    btn_openfolder = gr.Button('📂', scale=2)
                 with gr.Column():
-                    md_vram_table = gr.HTML(NLLB_VRAM_TABLE, elem_id="md_nllb_vram_table")
+                    gr.Markdown(MARKDOWN, elem_id="md_project")
+            with gr.Tabs():
+                with gr.TabItem("File"):  # tab1
+                    with gr.Row():
+                        input_file = gr.Files(type="file", label="Upload File here")
+                    with gr.Row():
+                        dd_model = gr.Dropdown(choices=self.whisper_inf.available_models, value="large-v2",
+                                               label="Model")
+                        dd_lang = gr.Dropdown(choices=["Automatic Detection"] + self.whisper_inf.available_langs,
+                                              value="Automatic Detection", label="Language")
+                        dd_subformat = gr.Dropdown(["SRT", "WebVTT"], value="SRT", label="Subtitle Format")
+                    with gr.Row():
+                        cb_translate = gr.Checkbox(value=False, label="Translate to English?", interactive=True)
+                    with gr.Row():
+                        btn_run = gr.Button("GENERATE SUBTITLE FILE", variant="primary")
+                    with gr.Row():
+                        tb_indicator = gr.Textbox(label="Output", scale=8)
+                        btn_openfolder = gr.Button('📂', scale=2)
+
+                    btn_run.click(fn=self.whisper_inf.transcribe_file,
+                                  inputs=[input_file, dd_model, dd_lang, dd_subformat, cb_translate],
+                                  outputs=[tb_indicator])
+                    btn_openfolder.click(fn=lambda: self.open_folder("outputs"), inputs=None, outputs=None)
+                    dd_model.change(fn=self.on_change_models, inputs=[dd_model], outputs=[cb_translate])
+
+                with gr.TabItem("Youtube"):  # tab2
+                    with gr.Row():
+                        tb_youtubelink = gr.Textbox(label="Youtube Link")
+                    with gr.Row(equal_height=True):
+                        with gr.Column():
+                            img_thumbnail = gr.Image(label="Youtube Thumbnail")
+                        with gr.Column():
+                            tb_title = gr.Label(label="Youtube Title")
+                            tb_description = gr.Textbox(label="Youtube Description", max_lines=15)
+                    with gr.Row():
+                        dd_model = gr.Dropdown(choices=self.whisper_inf.available_models, value="large-v2",
+                                               label="Model")
+                        dd_lang = gr.Dropdown(choices=["Automatic Detection"] + self.whisper_inf.available_langs,
+                                              value="Automatic Detection", label="Language")
+                        dd_subformat = gr.Dropdown(choices=["SRT", "WebVTT"], value="SRT", label="Subtitle Format")
+                    with gr.Row():
+                        cb_translate = gr.Checkbox(value=False, label="Translate to English?", interactive=True)
+                    with gr.Row():
+                        btn_run = gr.Button("GENERATE SUBTITLE FILE", variant="primary")
+                    with gr.Row():
+                        tb_indicator = gr.Textbox(label="Output", scale=8)
+                        btn_openfolder = gr.Button('📂', scale=2)
+
+                    btn_run.click(fn=self.whisper_inf.transcribe_youtube,
+                                  inputs=[tb_youtubelink, dd_model, dd_lang, dd_subformat, cb_translate],
+                                  outputs=[tb_indicator])
+                    tb_youtubelink.change(get_ytmetas, inputs=[tb_youtubelink],
+                                          outputs=[img_thumbnail, tb_title, tb_description])
+                    btn_openfolder.click(fn=lambda: self.open_folder("outputs"), inputs=None, outputs=None)
+                    dd_model.change(fn=self.on_change_models, inputs=[dd_model], outputs=[cb_translate])
+
+                with gr.TabItem("Mic"):  # tab3
+                    with gr.Row():
+                        mic_input = gr.Microphone(label="Record with Mic", type="filepath", interactive=True)
+                    with gr.Row():
+                        dd_model = gr.Dropdown(choices=self.whisper_inf.available_models, value="large-v2",
+                                               label="Model")
+                        dd_lang = gr.Dropdown(choices=["Automatic Detection"] + self.whisper_inf.available_langs,
+                                              value="Automatic Detection", label="Language")
+                        dd_subformat = gr.Dropdown(["SRT", "WebVTT"], value="SRT", label="Subtitle Format")
+                    with gr.Row():
+                        cb_translate = gr.Checkbox(value=False, label="Translate to English?", interactive=True)
+                    with gr.Row():
+                        btn_run = gr.Button("GENERATE SUBTITLE FILE", variant="primary")
+                    with gr.Row():
+                        tb_indicator = gr.Textbox(label="Output", scale=8)
+                        btn_openfolder = gr.Button('📂', scale=2)
+
+                    btn_run.click(fn=self.whisper_inf.transcribe_mic,
+                                  inputs=[mic_input, dd_model, dd_lang, dd_subformat, cb_translate],
+                                  outputs=[tb_indicator])
+                    btn_openfolder.click(fn=lambda: self.open_folder("outputs"), inputs=None, outputs=None)
+                    dd_model.change(fn=self.on_change_models, inputs=[dd_model], outputs=[cb_translate])
+
+                with gr.TabItem("T2T Translation"):  # tab 4
+                    with gr.Row():
+                        file_subs = gr.Files(type="file", label="Upload Subtitle Files to translate here",
+                                             file_types=['.vtt', '.srt'])
+
+                    with gr.TabItem("NLLB"):  # sub tab1
+                        with gr.Row():
+                            dd_nllb_model = gr.Dropdown(label="Model", value=self.nllb_inf.default_model_size,
+                                                        choices=self.nllb_inf.available_models)
+                            dd_nllb_sourcelang = gr.Dropdown(label="Source Language",
+                                                             choices=self.nllb_inf.available_source_langs)
+                            dd_nllb_targetlang = gr.Dropdown(label="Target Language",
+                                                             choices=self.nllb_inf.available_target_langs)
+                        with gr.Row():
+                            btn_run = gr.Button("TRANSLATE SUBTITLE FILE", variant="primary")
+                        with gr.Row():
+                            tb_indicator = gr.Textbox(label="Output", scale=8)
+                            btn_openfolder = gr.Button('📂', scale=2)
+                        with gr.Column():
+                            md_vram_table = gr.HTML(NLLB_VRAM_TABLE, elem_id="md_nllb_vram_table")
+
+                    btn_run.click(fn=self.nllb_inf.translate_file,
+                                  inputs=[file_subs, dd_nllb_model, dd_nllb_sourcelang, dd_nllb_targetlang],
+                                  outputs=[tb_indicator])
+                    btn_openfolder.click(fn=lambda: self.open_folder(os.path.join("outputs", "translations")),
+                                         inputs=None,
+                                         outputs=None)
+
+        if self.args.share:
+            self.app.queue(api_open=False).launch(share=True)
+        else:
+            self.app.queue(api_open=False).launch()
 
-            btn_run.click(fn=nllb_inf.translate_file,
-                          inputs=[file_subs, dd_nllb_model, dd_nllb_sourcelang, dd_nllb_targetlang],
-                          outputs=[tb_indicator])
-            btn_openfolder.click(fn=lambda: open_folder(os.path.join("outputs", "translations")), inputs=None, outputs=None)
 
+# Create the parser
+parser = argparse.ArgumentParser()
+parser.add_argument('--share', type=bool, default=False, nargs='?', const=True,
+                    help='Share value')
+_args = parser.parse_args()
 
-if args.share:
-    block.launch(share=True)
-else:
-    block.launch()
+if __name__ == "__main__":
+    app = App(args=_args)
+    app.launch()

modules/nllb_inference.py

@@ -1,10 +1,10 @@
-from .base_interface import BaseInterface
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 import gradio as gr
 import torch
 import os
 from datetime import datetime
 
+from .base_interface import BaseInterface
 from modules.subtitle_manager import *
 
 DEFAULT_MODEL_SIZE = "facebook/nllb-200-1.3B"
@@ -28,9 +28,29 @@ class NLLBInference(BaseInterface):
         result = self.pipeline(text)
         return result[0]['translation_text']
 
-    def translate_file(self, fileobjs
-                       , model_size, src_lang, tgt_lang,
+    def translate_file(self,
+                       fileobjs: list,
+                       model_size: str,
+                       src_lang: str,
+                       tgt_lang: str,
                        progress=gr.Progress()):
+        """
+        Translate subtitle file from source language to target language
+
+        Parameters
+        ----------
+        fileobjs: list
+            List of files to transcribe from gr.Files()
+        model_size: str
+            Whisper model size from gr.Dropdown()
+        src_lang: str
+            Source language of the file to translate from gr.Dropdown()
+        tgt_lang: str
+            Target language of the file to translate from gr.Dropdown()
+        progress: gr.Progress
+            Indicator to show progress directly in gradio.
+            I use a forked version of whisper for this. To see more info : https://github.com/jhj0517/jhj0517-whisper/tree/add-progress-callback
+        """
         try:
             if model_size != self.current_model_size or self.model is None:
                 print("\nInitializing NLLB Model..\n")

modules/whisper_Inference.py

@@ -1,11 +1,12 @@
 import whisper
-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
 import gradio as gr
 import os
 from datetime import datetime
 
+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
+
 DEFAULT_MODEL_SIZE = "large-v2"
 
 
@@ -17,10 +18,33 @@ class WhisperInference(BaseInterface):
         self.available_models = whisper.available_models()
         self.available_langs = sorted(list(whisper.tokenizer.LANGUAGES.values()))
 
-    def transcribe_file(self, fileobjs,
-                        model_size, lang, subformat, istranslate,
+    def transcribe_file(self,
+                        fileobjs: list,
+                        model_size: str,
+                        lang: str,
+                        subformat: str,
+                        istranslate: bool,
                         progress=gr.Progress()):
-
+        """
+        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.
+        progress: gr.Progress
+            Indicator to show progress directly in gradio.
+            I use a forked version of whisper for this. To see more info : https://github.com/jhj0517/jhj0517-whisper/tree/add-progress-callback
+        """
         def progress_callback(progress_value):
             progress(progress_value, desc="Transcribing..")
 
@@ -78,10 +102,33 @@ class WhisperInference(BaseInterface):
             self.release_cuda_memory()
             self.remove_input_files([fileobj.name for fileobj in fileobjs])
 
-    def transcribe_youtube(self, youtubelink,
-                           model_size, lang, subformat, istranslate,
+    def transcribe_youtube(self,
+                           youtubelink: str,
+                           model_size: str,
+                           lang: str,
+                           subformat: str,
+                           istranslate: bool,
                            progress=gr.Progress()):
-
+        """
+        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.
+        progress: gr.Progress
+            Indicator to show progress directly in gradio.
+            I use a forked version of whisper for this. To see more info : https://github.com/jhj0517/jhj0517-whisper/tree/add-progress-callback
+        """
         def progress_callback(progress_value):
             progress(progress_value, desc="Transcribing..")
 
@@ -128,10 +175,33 @@ class WhisperInference(BaseInterface):
             self.release_cuda_memory()
             self.remove_input_files([file_path])
 
-    def transcribe_mic(self, micaudio,
-                       model_size, lang, subformat, istranslate,
+    def transcribe_mic(self,
+                       micaudio: str,
+                       model_size: str,
+                       lang: str,
+                       subformat: str,
+                       istranslate: bool,
                        progress=gr.Progress()):
-
+        """
+        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.
+            I use a forked version of whisper for this. To see more info : https://github.com/jhj0517/jhj0517-whisper/tree/add-progress-callback
+        """
         def progress_callback(progress_value):
             progress(progress_value, desc="Transcribing..")