Merge pull request #223 from jhj0517/refactor/abstract-class

app.py

@@ -257,6 +257,9 @@ class App:
                                                                   self.deepl_api.available_target_langs.keys()))
                         with gr.Row():
                             cb_deepl_ispro = gr.Checkbox(label="Pro User?", value=False)
+                        with gr.Row():
+                            cb_timestamp = gr.Checkbox(value=True, label="Add a timestamp to the end of the filename",
+                                                       interactive=True)
                         with gr.Row():
                             btn_run = gr.Button("TRANSLATE SUBTITLE FILE", variant="primary")
                         with gr.Row():
@@ -266,7 +269,7 @@ class App:
 
                     btn_run.click(fn=self.deepl_api.translate_deepl,
                                   inputs=[tb_authkey, file_subs, dd_deepl_sourcelang, dd_deepl_targetlang,
-                                          cb_deepl_ispro],
+                                          cb_deepl_ispro, cb_timestamp],
                                   outputs=[tb_indicator, files_subtitles])
 
                     btn_openfolder.click(fn=lambda: self.open_folder(os.path.join("outputs", "translations")),

modules/translation/deepl_api.py

@@ -83,7 +83,7 @@ DEEPL_AVAILABLE_SOURCE_LANGS = {
 
 class DeepLAPI:
     def __init__(self,
-                 output_dir: str
+                 output_dir: str = os.path.join("outputs", "translations")
                  ):
         self.api_interval = 1
         self.max_text_batch_size = 50
@@ -97,6 +97,7 @@ class DeepLAPI:
                         source_lang: str,
                         target_lang: str,
                         is_pro: bool,
+                        add_timestamp: bool,
                         progress=gr.Progress()) -> list:
         """
         Translate subtitle files using DeepL API
@@ -112,6 +113,8 @@ class DeepLAPI:
             Target language of the file to transcribe from gr.Dropdown()
         is_pro: str
             Boolean value that is about pro user or not from gr.Checkbox().
+        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.
 
@@ -141,10 +144,12 @@ class DeepLAPI:
                     progress(batch_end / len(parsed_dicts), desc="Translating..")
 
                 subtitle = get_serialized_srt(parsed_dicts)
-                timestamp = datetime.now().strftime("%m%d%H%M%S")
 
-                file_name = file_name[:-9]
-                output_path = os.path.join(self.output_dir, "", f"{file_name}-{timestamp}.srt")
+                if add_timestamp:
+                    timestamp = datetime.now().strftime("%m%d%H%M%S")
+                    file_name += f"-{timestamp}"
+
+                output_path = os.path.join(self.output_dir, "", f"{file_name}.srt")
                 write_file(subtitle, output_path)
 
             elif file_ext == ".vtt":
@@ -161,11 +166,12 @@ class DeepLAPI:
                     progress(batch_end / len(parsed_dicts), desc="Translating..")
 
                 subtitle = get_serialized_vtt(parsed_dicts)
-                timestamp = datetime.now().strftime("%m%d%H%M%S")
 
-                file_name = file_name[:-9]
-                output_path = os.path.join(self.output_dir, "", f"{file_name}-{timestamp}.vtt")
+                if add_timestamp:
+                    timestamp = datetime.now().strftime("%m%d%H%M%S")
+                    file_name += f"-{timestamp}"
 
+                output_path = os.path.join(self.output_dir, "", f"{file_name}.vtt")
                 write_file(subtitle, output_path)
 
             files_info[file_name] = {"subtitle": subtitle, "path": output_path}

modules/translation/nllb_inference.py

@@ -7,8 +7,8 @@ from modules.translation.translation_base import TranslationBase
 
 class NLLBInference(TranslationBase):
     def __init__(self,
-                 model_dir: str,
-                 output_dir: str
+                 model_dir: str = os.path.join("models", "NLLB"),
+                 output_dir: str = os.path.join("outputs", "translations")
                  ):
         super().__init__(
             model_dir=model_dir,

modules/translation/translation_base.py

@@ -11,8 +11,9 @@ from modules.utils.subtitle_manager import *
 
 class TranslationBase(ABC):
     def __init__(self,
-                 model_dir: str,
-                 output_dir: str):
+                 model_dir: str = os.path.join("models", "NLLB"),
+                 output_dir: str = os.path.join("outputs", "translations")
+                 ):
         super().__init__()
         self.model = None
         self.model_dir = model_dir

modules/whisper/faster_whisper_inference.py

@@ -17,17 +17,15 @@ from modules.whisper.whisper_base import WhisperBase
 
 class FasterWhisperInference(WhisperBase):
     def __init__(self,
-                 model_dir: Optional[str] = None,
-                 diarization_model_dir: Optional[str] = None,
-                 output_dir: Optional[str] = None,
+                 model_dir: str = os.path.join("models", "Whisper", "faster-whisper"),
+                 diarization_model_dir: str = os.path.join("models", "Diarization"),
+                 output_dir: str = os.path.join("outputs"),
                  ):
         super().__init__(
             model_dir=model_dir,
             diarization_model_dir=diarization_model_dir,
             output_dir=output_dir
         )
-        if model_dir is None:
-            model_dir = os.path.join("models", "Whisper", "faster-whisper")
         self.model_dir = model_dir
         os.makedirs(self.model_dir, exist_ok=True)
 

modules/whisper/insanely_fast_whisper_inference.py

@@ -17,17 +17,15 @@ from modules.whisper.whisper_base import WhisperBase
 
 class InsanelyFastWhisperInference(WhisperBase):
     def __init__(self,
-                 model_dir: Optional[str] = None,
-                 diarization_model_dir: Optional[str] = None,
-                 output_dir: Optional[str] = None,
+                 model_dir: str = os.path.join("models", "Whisper", "insanely-fast-whisper"),
+                 diarization_model_dir: str = os.path.join("models", "Diarization"),
+                 output_dir: str = os.path.join("outputs"),
                  ):
         super().__init__(
             model_dir=model_dir,
             output_dir=output_dir,
             diarization_model_dir=diarization_model_dir
         )
-        if model_dir is None:
-            model_dir = os.path.join("models", "Whisper", "insanely-fast-whisper")
         self.model_dir = model_dir
         os.makedirs(self.model_dir, exist_ok=True)
 

modules/whisper/whisper_Inference.py

@@ -4,6 +4,7 @@ import time
 from typing import BinaryIO, Union, Tuple, List
 import numpy as np
 import torch
+import os
 from argparse import Namespace
 
 from modules.whisper.whisper_base import WhisperBase
@@ -12,9 +13,9 @@ from modules.whisper.whisper_parameter import *
 
 class WhisperInference(WhisperBase):
     def __init__(self,
-                 model_dir: Optional[str] = None,
-                 diarization_model_dir: Optional[str] = None,
-                 output_dir: Optional[str] = None,
+                 model_dir: str = os.path.join("models", "Whisper"),
+                 diarization_model_dir: str = os.path.join("models", "Diarization"),
+                 output_dir: str = os.path.join("outputs"),
                  ):
         super().__init__(
             model_dir=model_dir,

modules/whisper/whisper_base.py

@@ -19,17 +19,10 @@ from modules.vad.silero_vad import SileroVAD
 
 class WhisperBase(ABC):
     def __init__(self,
-                 model_dir: Optional[str] = None,
-                 diarization_model_dir: Optional[str] = None,
-                 output_dir: Optional[str] = None,
+                 model_dir: str = os.path.join("models", "Whisper"),
+                 diarization_model_dir: str = os.path.join("models", "Diarization"),
+                 output_dir: str = os.path.join("outputs"),
                  ):
-        if model_dir is None:
-            model_dir = os.path.join("models", "Whisper")
-        if diarization_model_dir is None:
-            diarization_model_dir = os.path.join("models", "Diarization")
-        if output_dir is None:
-            output_dir = os.path.join("outputs")
-
         self.model_dir = model_dir
         self.output_dir = output_dir
         os.makedirs(self.output_dir, exist_ok=True)

modules/whisper/whisper_factory.py

@@ -11,11 +11,11 @@ class WhisperFactory:
     @staticmethod
     def create_whisper_inference(
         whisper_type: str,
-        whisper_model_dir: Optional[str] = None,
-        faster_whisper_model_dir: Optional[str] = None,
-        insanely_fast_whisper_model_dir: Optional[str] = None,
-        diarization_model_dir: Optional[str] = None,
-        output_dir: Optional[str] = None,
+        whisper_model_dir: str = os.path.join("models", "Whisper"),
+        faster_whisper_model_dir: str = os.path.join("models", "Whisper", "faster-whisper"),
+        insanely_fast_whisper_model_dir: str = os.path.join("models", "Whisper", "insanely-fast-whisper"),
+        diarization_model_dir: str = os.path.join("models", "Diarization"),
+        output_dir: str = os.path.join("outputs"),
     ) -> "WhisperBase":
         """
         Create a whisper inference class based on the provided whisper_type.