Upload 4 files

modules/translation/deepl_api.py

@@ -5,6 +5,7 @@ from datetime import datetime
 import gradio as gr
 
 from modules.utils.paths import TRANSLATION_OUTPUT_DIR, DEFAULT_PARAMETERS_CONFIG_PATH
+from modules.utils.constants import AUTOMATIC_DETECTION
 from modules.utils.subtitle_manager import *
 from modules.utils.files_manager import load_yaml, save_yaml
 
@@ -50,7 +51,7 @@ DEEPL_AVAILABLE_TARGET_LANGS = {
 }
 
 DEEPL_AVAILABLE_SOURCE_LANGS = {
-    'Automatic Detection': None,
+    AUTOMATIC_DETECTION: None,
     'Bulgarian': 'BG',
     'Czech': 'CS',
     'Danish': 'DA',
@@ -138,37 +139,27 @@ class DeepLAPI:
         )
 
         files_info = {}
-        for fileobj in fileobjs:
-            file_path = fileobj
-            file_name, file_ext = os.path.splitext(os.path.basename(fileobj))
-
-            if file_ext == ".srt":
-                parsed_dicts = parse_srt(file_path=file_path)
-
-            elif file_ext == ".vtt":
-                parsed_dicts = parse_vtt(file_path=file_path)
+        for file_path in fileobjs:
+            file_name, file_ext = os.path.splitext(os.path.basename(file_path))
+            writer = get_writer(file_ext, self.output_dir)
+            segments = writer.to_segments(file_path)
 
             batch_size = self.max_text_batch_size
-            for batch_start in range(0, len(parsed_dicts), batch_size):
-                batch_end = min(batch_start + batch_size, len(parsed_dicts))
-                sentences_to_translate = [dic["sentence"] for dic in parsed_dicts[batch_start:batch_end]]
+            for batch_start in range(0, len(segments), batch_size):
+                progress(batch_start / len(segments), desc="Translating..")
+                sentences_to_translate = [seg.text for seg in segments[batch_start:batch_start+batch_size]]
                 translated_texts = self.request_deepl_translate(auth_key, sentences_to_translate, source_lang,
                                                                 target_lang, is_pro)
                 for i, translated_text in enumerate(translated_texts):
-                    parsed_dicts[batch_start + i]["sentence"] = translated_text["text"]
-                progress(batch_end / len(parsed_dicts), desc="Translating..")
-
-            if file_ext == ".srt":
-                subtitle = get_serialized_srt(parsed_dicts)
-            elif file_ext == ".vtt":
-                subtitle = get_serialized_vtt(parsed_dicts)
-
-            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}{file_ext}")
-            write_file(subtitle, output_path)
+                    segments[batch_start + i].text = translated_text["text"]
+
+            subtitle, output_path = generate_file(
+                output_dir=self.output_dir,
+                output_file_name=file_name,
+                output_format=file_ext,
+                result=segments,
+                add_timestamp=add_timestamp
+            )
 
             files_info[file_name] = {"subtitle": subtitle, "path": output_path}
 

modules/translation/nllb_inference.py

@@ -3,10 +3,10 @@ import gradio as gr
 import os
 
 from modules.utils.paths import TRANSLATION_OUTPUT_DIR, NLLB_MODELS_DIR
-from modules.translation.translation_base import TranslationBase
+import modules.translation.translation_base as base
 
 
-class NLLBInference(TranslationBase):
+class NLLBInference(base.TranslationBase):
     def __init__(self,
                  model_dir: str = NLLB_MODELS_DIR,
                  output_dir: str = TRANSLATION_OUTPUT_DIR
@@ -29,7 +29,7 @@ class NLLBInference(TranslationBase):
             text,
             max_length=max_length
         )
-        return result[0]['translation_text']
+        return result[0]["translation_text"]
 
     def update_model(self,
                      model_size: str,
@@ -41,8 +41,7 @@ class NLLBInference(TranslationBase):
             if lang in NLLB_AVAILABLE_LANGS:
                 return NLLB_AVAILABLE_LANGS[lang]
             elif lang not in NLLB_AVAILABLE_LANGS.values():
-                raise ValueError(
-                    f"Language '{lang}' is not supported. Use one of: {list(NLLB_AVAILABLE_LANGS.keys())}")
+                raise ValueError(f"Language '{lang}' is not supported. Use one of: {list(NLLB_AVAILABLE_LANGS.keys())}")
             return lang
 
         src_lang = validate_language(src_lang)

modules/translation/translation_base.py

@@ -2,14 +2,17 @@ import os
 import torch
 import gradio as gr
 from abc import ABC, abstractmethod
+import gc
 from typing import List
 from datetime import datetime
 
-from modules.whisper.whisper_parameter import *
+import modules.translation.nllb_inference as nllb
+from modules.whisper.data_classes import *
 from modules.utils.subtitle_manager import *
 from modules.utils.files_manager import load_yaml, save_yaml
 from modules.utils.paths import DEFAULT_PARAMETERS_CONFIG_PATH, NLLB_MODELS_DIR, TRANSLATION_OUTPUT_DIR
 
+
 class TranslationBase(ABC):
     def __init__(self,
                  model_dir: str = NLLB_MODELS_DIR,
@@ -93,32 +96,22 @@ class TranslationBase(ABC):
             files_info = {}
             for fileobj in fileobjs:
                 file_name, file_ext = os.path.splitext(os.path.basename(fileobj))
-                if file_ext == ".srt":
-                    parsed_dicts = parse_srt(file_path=fileobj)
-                    total_progress = len(parsed_dicts)
-                    for index, dic in enumerate(parsed_dicts):
-                        progress(index / total_progress, desc="Translating..")
-                        translated_text = self.translate(dic["sentence"], max_length=max_length)
-                        dic["sentence"] = translated_text
-                    subtitle = get_serialized_srt(parsed_dicts)
-
-                elif file_ext == ".vtt":
-                    parsed_dicts = parse_vtt(file_path=fileobj)
-                    total_progress = len(parsed_dicts)
-                    for index, dic in enumerate(parsed_dicts):
-                        progress(index / total_progress, desc="Translating..")
-                        translated_text = self.translate(dic["sentence"], max_length=max_length)
-                        dic["sentence"] = translated_text
-                    subtitle = get_serialized_vtt(parsed_dicts)
-
-                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}{file_ext}")
-                write_file(subtitle, output_path)
-
-                files_info[file_name] = {"subtitle": subtitle, "path": output_path}
+                writer = get_writer(file_ext, self.output_dir)
+                segments = writer.to_segments(fileobj)
+                for i, segment in enumerate(segments):
+                    progress(i / len(segments), desc="Translating..")
+                    translated_text = self.translate(segment.text, max_length=max_length)
+                    segment.text = translated_text
+
+                subtitle, file_path = generate_file(
+                    output_dir=self.output_dir,
+                    output_file_name=file_name,
+                    output_format=file_ext,
+                    result=segments,
+                    add_timestamp=add_timestamp
+                )
+
+                files_info[file_name] = {"subtitle": subtitle, "path": file_path}
 
             total_result = ''
             for file_name, info in files_info.items():
@@ -131,61 +124,20 @@ class TranslationBase(ABC):
             return [gr_str, output_file_paths]
 
         except Exception as e:
-            print(f"Error: {str(e)}")
+            print(f"Error translating file: {e}")
+            raise
         finally:
             self.release_cuda_memory()
 
-    def translate_text(self,
-                       input_list_dict: list,
-                       model_size: str,
-                       src_lang: str,
-                       tgt_lang: str,
-                       max_length: int = 200,
-                       add_timestamp: bool = True,
-                       progress=gr.Progress()) -> list:
-        """
-        Translate text from source language to target language
-
-        Parameters
-        ----------
-        str_text: str
-            List[dict] to translate
-        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()
-        max_length: int
-            Max length per line to translate
-        add_timestamp: bool
-            Boolean value that determines whether to add a timestamp
-        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
-
-        Returns
-        ----------
-        A List of
-        List[dict] with translation
-        """
-        try:
-            self.cache_parameters(model_size=model_size,src_lang=src_lang,tgt_lang=tgt_lang,max_length=max_length,add_timestamp=add_timestamp)
-            self.update_model(model_size=model_size,src_lang=src_lang,tgt_lang=tgt_lang,progress=progress)
-            
-            total_progress = len(input_list_dict)
-            for index, dic in enumerate(input_list_dict):
-                progress(index / total_progress, desc="Translating..")
-                translated_text = self.translate(dic["text"], max_length=max_length)
-                dic["text"] = translated_text
-
-            return input_list_dict
-
-        except Exception as e:
-            print(f"Error: {str(e)}")
-        finally:
+    def offload(self):
+        """Offload the model and free up the memory"""
+        if self.model is not None:
+            del self.model
+            self.model = None
+        if self.device == "cuda":
             self.release_cuda_memory()
-    
+        gc.collect()
+
     @staticmethod
     def get_device():
         if torch.cuda.is_available():
@@ -216,11 +168,17 @@ class TranslationBase(ABC):
                          tgt_lang: str,
                          max_length: int,
                          add_timestamp: bool):
+        def validate_lang(lang: str):
+            if lang in list(nllb.NLLB_AVAILABLE_LANGS.values()):
+                flipped = {value: key for key, value in nllb.NLLB_AVAILABLE_LANGS.items()}
+                return flipped[lang]
+            return lang
+
         cached_params = load_yaml(DEFAULT_PARAMETERS_CONFIG_PATH)
         cached_params["translation"]["nllb"] = {
             "model_size": model_size,
-            "source_lang": src_lang,
-            "target_lang": tgt_lang,
+            "source_lang": validate_lang(src_lang),
+            "target_lang": validate_lang(tgt_lang),
             "max_length": max_length,
         }
         cached_params["translation"]["add_timestamp"] = add_timestamp