add output_dir

modules/deepl_api.py

@@ -82,11 +82,14 @@ DEEPL_AVAILABLE_SOURCE_LANGS = {
 
 
 class DeepLAPI:
-    def __init__(self):
+    def __init__(self,
+                 output_dir: str
+                 ):
         self.api_interval = 1
         self.max_text_batch_size = 50
         self.available_target_langs = DEEPL_AVAILABLE_TARGET_LANGS
         self.available_source_langs = DEEPL_AVAILABLE_SOURCE_LANGS
+        self.output_dir = output_dir
 
     def translate_deepl(self,
                         auth_key: str,
@@ -111,6 +114,7 @@ class DeepLAPI:
             Boolean value that is about pro user or not from gr.Checkbox().
         progress: gr.Progress
             Indicator to show progress directly in gradio.
+
         Returns
         ----------
         A List of
@@ -140,7 +144,7 @@ class DeepLAPI:
                 timestamp = datetime.now().strftime("%m%d%H%M%S")
 
                 file_name = file_name[:-9]
-                output_path = os.path.join("outputs", "translations", f"{file_name}-{timestamp}.srt")
+                output_path = os.path.join(self.output_dir, "translations", f"{file_name}-{timestamp}.srt")
                 write_file(subtitle, output_path)
 
             elif file_ext == ".vtt":
@@ -160,7 +164,7 @@ class DeepLAPI:
                 timestamp = datetime.now().strftime("%m%d%H%M%S")
 
                 file_name = file_name[:-9]
-                output_path = os.path.join("outputs", "translations", f"{file_name}-{timestamp}.vtt")
+                output_path = os.path.join(self.output_dir, "translations", f"{file_name}-{timestamp}.vtt")
 
                 write_file(subtitle, output_path)
 

modules/nllb_inference.py

@@ -6,9 +6,13 @@ from modules.translation_base import TranslationBase
 
 
 class NLLBInference(TranslationBase):
-    def __init__(self):
+    def __init__(self,
+                 model_dir: str,
+                 output_dir: str
+                 ):
         super().__init__(
-            model_dir=os.path.join("models", "NLLB")
+            model_dir=model_dir,
+            output_dir=output_dir
         )
         self.tokenizer = None
         self.available_models = ["facebook/nllb-200-3.3B", "facebook/nllb-200-1.3B", "facebook/nllb-200-distilled-600M"]

modules/translation_base.py

@@ -11,11 +11,14 @@ from modules.subtitle_manager import *
 
 class TranslationBase(ABC):
     def __init__(self,
-                 model_dir: str):
+                 model_dir: str,
+                 output_dir: str):
         super().__init__()
         self.model = None
         self.model_dir = model_dir
+        self.output_dir = output_dir
         os.makedirs(self.model_dir, exist_ok=True)
+        os.makedirs(self.output_dir, exist_ok=True)
         self.current_model_size = None
         self.device = self.get_device()
 
@@ -102,9 +105,9 @@ class TranslationBase(ABC):
 
                     timestamp = datetime.now().strftime("%m%d%H%M%S")
                     if add_timestamp:
-                        output_path = os.path.join("outputs", "translations", f"{file_name}-{timestamp}")
+                        output_path = os.path.join(self.output_dir, "translations", f"{file_name}-{timestamp}")
                     else:
-                        output_path = os.path.join("outputs", "translations", f"{file_name}.vtt")
+                        output_path = os.path.join(self.output_dir, "translations", f"{file_name}.vtt")
 
                 write_file(subtitle, output_path)
                 files_info[file_name] = subtitle