Upload 4 files

modules/translation/deepl_api.py

@@ -5,7 +5,6 @@ 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
 
@@ -51,7 +50,7 @@ DEEPL_AVAILABLE_TARGET_LANGS = {
 }
 
 DEEPL_AVAILABLE_SOURCE_LANGS = {
-    AUTOMATIC_DETECTION: None,
+    'Automatic Detection': None,
     'Bulgarian': 'BG',
     'Czech': 'CS',
     'Danish': 'DA',
@@ -139,27 +138,37 @@ class DeepLAPI:
         )
 
         files_info = {}
-        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)
+        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)
 
             batch_size = self.max_text_batch_size
-            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]]
+            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]]
                 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):
-                    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
-            )
+                    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)
 
             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
-import modules.translation.translation_base as base
+from modules.translation.translation_base import TranslationBase
 
 
-class NLLBInference(base.TranslationBase):
+class NLLBInference(TranslationBase):
     def __init__(self,
                  model_dir: str = NLLB_MODELS_DIR,
                  output_dir: str = TRANSLATION_OUTPUT_DIR
@@ -29,7 +29,7 @@ class NLLBInference(base.TranslationBase):
             text,
             max_length=max_length
         )
-        return result[0]["translation_text"]
+        return result[0]['translation_text']
 
     def update_model(self,
                      model_size: str,
@@ -41,7 +41,8 @@ class NLLBInference(base.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,11 +2,10 @@ import os
 import torch
 import gradio as gr
 from abc import ABC, abstractmethod
-import gc
 from typing import List
 from datetime import datetime
 
-import modules.translation.nllb_inference as nllb
+from modules.whisper.whisper_parameter 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
@@ -95,22 +94,32 @@ class TranslationBase(ABC):
             files_info = {}
             for fileobj in fileobjs:
                 file_name, file_ext = os.path.splitext(os.path.basename(fileobj))
-                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}
+                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}
 
             total_result = ''
             for file_name, info in files_info.items():
@@ -123,20 +132,10 @@ class TranslationBase(ABC):
             return [gr_str, output_file_paths]
 
         except Exception as e:
-            print(f"Error translating file: {e}")
-            raise
+            print(f"Error: {str(e)}")
         finally:
             self.release_cuda_memory()
 
-    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():
@@ -167,17 +166,11 @@ 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": validate_lang(src_lang),
-            "target_lang": validate_lang(tgt_lang),
+            "source_lang": src_lang,
+            "target_lang": tgt_lang,
             "max_length": max_length,
         }
         cached_params["translation"]["add_timestamp"] = add_timestamp