nemo

Co-authored-by: APodoinikov <APodoynikov@detmir.ru>

Dockerfile

@@ -4,7 +4,7 @@ WORKDIR /app
 
 ENV PYTHONUNBUFFERED=1
 
-RUN apt update && apt install -y python3-venv python3-pip
+RUN apt update && apt install -y python3-venv python3-pip ffmpeg
 
 COPY requirements.txt .
 
@@ -13,7 +13,6 @@ RUN ./venv/bin/pip3 install --no-cache-dir -r requirements.txt
 
 COPY jaa.py ./
 COPY main.py ./
-COPY log_config.yaml ./
 COPY app ./app
 COPY plugins ./plugins
 COPY static ./static

app/app_core.py

@@ -248,9 +248,12 @@ class AppCore(JaaCore):
             return dto.ProcessingFileDirResp(files=list(), error=getattr(e, 'message', repr(e)))
         finally:
             # clear memory (GPU or RAM) for heavy file processing plugins
-            for processors_list in self.files_ext_to_processors.values():
-                for processor in processors_list:
-                    processor.after_processing_function(self)
+            try:
+                for processors_list in self.files_ext_to_processors.values():
+                    for processor in processors_list:
+                        processor.after_processing_function(self)
+            except Exception as fe:
+                log.log_exception("Error after processing function: ", fe)
 
     def process_file(self, req: dto.ProcessingFileDirReq, root: str, file_name: str) -> dto.ProcessingFileResp:
         try:

doc/ru/plugins-file-processing/file_media_nemo.md

@@ -0,0 +1,23 @@
+# Плагин обработки файлов: media
+
+Часть параметров, общих для всех плагинов, описана [здесь](../processing_files.md).
+
+## Параметры плагина
+
+* **model** - модель для распознавания звука.
+В данный момент основные модели - маленькая `nvidia/parakeet-tdt-0.6b-v3`, и большая `nvidia/canary-1b-v2`.
+
+* **cuda** - `true` - использовать видеокарту (быстрее), `false` - использовать cpu (медленнее).
+
+* **cuda_device_index** - если в системе несколько видеокарт, можно выбрать ту, в которую будет загружена модель.
+  Номер и имя видеокарты указывается при старте приложения в логе, вида `INFO GPU #0: NVIDIA GeForce RTX 4090`.
+  `0` - указываемый в параметре номер.
+
+* **unload_model_after_processing** - выгружать модель из памяти после завершения обработки списка файлов. Можно использовать, если на карте мало памяти.
+На этой модели замечены проблемы с этим параметром - модель так и остается в памяти, очищает память только перезапуск приложения.
+
+* **translate_after_processing** - переводить субтитры сразу после распознавания. По умолчанию, будет запущен обработчик [srt-Файлов](file_srt.md).
+
+* **output_file_name_template** - шаблон для имени файла.
+
+* **batch_size** - размер пачки токенов на обработку, ускорение обработки за счет большего размера памяти.

doc/ru/processing_files.md

@@ -18,6 +18,10 @@
 * **file_txt** - [документация](plugins-file-processing/file_txt.md). Перевод книг в формате txt.
 
 * **file_media_whisper** - [документация](plugins-file-processing/file_media_whisper.md). Распознавание текста в медиа-файлах. Поддерживаются практически все файлы, которые может прочитать ffmpeg.
+* 
+* **file_media_nemo** - [документация](plugins-file-processing/file_media_nemo.md). Распознавание текста в медиа-файлах. Поддерживаются практически все файлы, которые может прочитать ffmpeg.
+Это альтернатива `file_media_whisper` - по умолчанию запускается Whisper, так как он поддерживает больше языков и более настраиваемый.
+Чтобы включить `file_media_nemo`, нужно выключить Whisper, или поменять в настройках типы файлов по умолчанию.
 
 * **file_srt** - [документация](plugins-file-processing/file_txt.md). Перевод субтитров srt. Можно найти существующие субтитры, или распознать звуковую дорожку через плагин `file_media_whisper`.
 

plugins/plugin_file_media_nemo.py

@@ -0,0 +1,157 @@
+import datetime
+import gc
+import inspect
+import os
+
+import torch
+from nemo.collections.asr.models import ASRModel, EncDecRNNTModel
+from pydub import AudioSegment
+
+from app import file_processor, cuda
+from app.app_core import AppCore
+from app.dto import ProcessingFileDirReq, ProcessingFileResp, FileProcessingPluginInitInfo, ProcessingFileStruct
+
+plugin_name = os.path.basename(__file__)[:-3]  # calculating modname
+
+model: EncDecRNNTModel | None = None
+
+
+def start(core: AppCore):
+    manifest = {  # plugin settings
+        "name": "Subtitle extractor for media files (Nemo)",  # name
+        "version": "1.0",  # version
+
+        "default_options": {
+            "enabled": True,
+            "model": "nvidia/canary-1b-v2",
+            "cuda": True,
+            "cuda_device_index": 0,
+            "unload_model_after_processing": True,
+            "translate_after_processing": True,
+
+            "batch_size": 4,
+
+            "output_file_name_template": "%%source%%.src_sub",
+
+            "default_extension_processor": {
+            },
+        },
+
+        "file_processing": {
+            "file_media_nemo_processing": (init, file_processing, processed_file_name, after_processing)
+        },
+    }
+
+    return manifest
+
+
+def start_with_options(core: AppCore, manifest: dict):
+    pass
+
+
+def init(core: AppCore) -> FileProcessingPluginInitInfo:
+    ext = {"mpeg", "mpg", "mp3", "mp4", "avi", "wav", "mkv", "vob", "ac3", "mpa", "ogg"}
+
+    return FileProcessingPluginInitInfo(plugin_name=plugin_name, supported_extensions=ext)
+
+
+def format_srt_time(seconds: float) -> str:
+    """Converts seconds to SRT time format HH:MM:SS,mmm using datetime.timedelta"""
+    sanitized_total_seconds = max(0.0, seconds)
+    delta = datetime.timedelta(seconds=sanitized_total_seconds)
+    total_int_seconds = int(delta.total_seconds())
+
+    hours = total_int_seconds // 3600
+    remainder_seconds_after_hours = total_int_seconds % 3600
+    minutes = remainder_seconds_after_hours // 60
+    seconds_part = remainder_seconds_after_hours % 60
+    milliseconds = delta.microseconds // 1000
+
+    return f"{hours:02d}:{minutes:02d}:{seconds_part:02d},{milliseconds:03d}"
+
+
+def generate_srt_content(segment_timestamps: list) -> str:
+    """Generates SRT formatted string from segment timestamps."""
+    srt_content = []
+    for i, ts in enumerate(segment_timestamps):
+        start_time = format_srt_time(ts['start'])
+        end_time = format_srt_time(ts['end'])
+        text = ts['segment']
+        srt_content.append(str(i + 1))
+        srt_content.append(f"{start_time} --> {end_time}")
+        srt_content.append(text)
+        srt_content.append("")
+    return "\n".join(srt_content)
+
+
+def file_processing(core: AppCore, file_struct: ProcessingFileStruct, req: ProcessingFileDirReq) -> ProcessingFileResp:
+    options = core.plugin_options(plugin_name)
+
+    global model
+    if model is None:
+        model = ASRModel.from_pretrained(model_name=options["model"])
+        model.to(cuda.get_device_with_gpu_num(options))
+
+    audio = AudioSegment.from_file(file_struct.path_file_in())
+    # supports only one channel and 16 000 frame rate
+    resampled_audio_file = f'{file_struct.path_out}{os.sep}{file_struct.file_name}_resampled.wav'
+    audio.export(resampled_audio_file, format="wav", parameters=["-ar", "16000", "-ac", "1"])
+
+    try:
+        # need to get function params - canary model supports "source_lang", parapet model - doesn't
+        param_names = [name for name, _ in inspect.signature(model.transcribe).parameters.items()]
+        if "source_lang" in param_names and "target_lang" in param_names:
+            transcribe = model.transcribe(audio=[resampled_audio_file], source_lang=req.from_lang, target_lang=req.from_lang,
+                                          timestamps=True, batch_size=options["batch_size"])
+        else:
+            transcribe = model.transcribe(audio=[resampled_audio_file],
+                                          timestamps=True, batch_size=options["batch_size"])
+        if not transcribe or not isinstance(transcribe, list) or not transcribe[0] or not hasattr(transcribe[0], 'timestamp') or not transcribe[0].timestamp or 'segment' not in transcribe[0].timestamp:
+            return file_processor.get_processing_file_resp_error(
+                file_in=file_struct.file_name_ext, path_in=file_struct.path_in, error_msg="Can't get transcribe")
+
+        segment_timestamps = transcribe[0].timestamp['segment'] # segment level timestamps
+        if segment_timestamps:
+            srt_content = generate_srt_content(segment_timestamps)
+
+            out_file_name = processed_file_name(core=core, file_struct=file_struct, req=req)
+            with open(file_struct.path_file_out(out_file_name), "w") as f:
+                f.write(srt_content)
+            if options["translate_after_processing"]:
+                return translate_after_processing(core=core, req=req, file_name_ext=out_file_name)
+            else:
+                return file_processor.get_processing_file_resp_ok(file_struct=file_struct, file_out=out_file_name)
+        else:
+            return file_processor.get_processing_file_resp_error(
+                file_in=file_struct.file_name_ext, path_in=file_struct.path_in, error_msg="Can't get segment timestamps")
+    finally:
+        if os.path.exists(resampled_audio_file):
+            os.remove(resampled_audio_file)
+
+
+def processed_file_name(core: AppCore, file_struct: ProcessingFileStruct, req: ProcessingFileDirReq) -> str:
+    options = core.plugin_options(plugin_name)
+    template: str = options["output_file_name_template"]
+
+    return file_processor.file_name_from_predefined_template(file_struct=file_struct, req=req,
+                                                             template=template, replace_ext="srt")
+
+
+def translate_after_processing(core: AppCore, req: ProcessingFileDirReq, file_name_ext: str) -> ProcessingFileResp:
+    return core.process_file(req=req, root=req.directory_out, file_name=file_name_ext)
+
+
+def after_processing(core: AppCore) -> None:
+    options = core.plugin_options(plugin_name)
+    global model
+
+    if options["unload_model_after_processing"] and model is not None:
+
+        model = None
+
+        if options["cuda"]:
+            torch.cuda.empty_cache()
+            del model
+        gc.collect()
+
+        model = None

plugins/plugin_file_media_whisper.py

@@ -38,25 +38,25 @@ def start(core: AppCore):
             "logprob_threshold": -1.0,
 
             "output_file_name_template": "%%source%%.src_sub",
+
+            "default_extension_processor": {
+                "mpeg": True,
+                "mpg": True,
+                "mp4": True,
+                "mp3": True,
+                "avi": True,
+                "wav": True,
+                "mkv": True,
+                "vob": True,
+                "ac3": True,
+                "mpa": True,
+                "ogg": True,
+            },
         },
 
         "file_processing": {
             "file_media_whisper_processing": (init, file_processing, processed_file_name, after_processing)
         },
-
-        "default_extension_processor": {
-            "mpeg": True,
-            "mpg": True,
-            "mp4": True,
-            "mp3": True,
-            "avi": True,
-            "wav": True,
-            "mkv": True,
-            "vob": True,
-            "ac3": True,
-            "mpa": True,
-            "ogg": True,
-        },
     }
 
     return manifest

requirements.txt

@@ -3,7 +3,9 @@ torch
 torchvision
 torchaudio
 openai-whisper == 20250625
-pysrt == 1.1.2
+nemo_toolkit[asr] == 2.5.0
+huggingface-hub == 0.35.3
+hf-xet == 1.1.10
 
 uvicorn == 0.34.2
 uvicorn[standard] == 0.34.2
@@ -12,6 +14,7 @@ termcolor == 3.1.0
 natsort == 8.4.0
 chardet == 5.2.0
 pyway == 0.3.32
+pysrt == 1.1.2
 
 transformers == 4.53.1
 ctranslate2 == 4.6.0

static/file_processing.js

@@ -54,6 +54,8 @@ async function process_files() {
     const toLang = document.getElementById('to_lang_select').value;
     const plugin = document.getElementById('plugin').value;
 
+    errorText.innerHTML = ""
+
     const reqBody = JSON.stringify({
         from_lang: fromLang, to_lang: toLang, translator_plugin: plugin,
         preserve_original_text: preserve_original_text, overwrite_processed_files: overwrite_processed_files,

static/index.js

@@ -12,6 +12,8 @@ async function translateText() {
     const toLang = document.getElementById('to_lang_select').value;
     const plugin = document.getElementById('plugin').value;
 
+    errorText.innerHTML = ""
+
     try {
         const reqBody = JSON.stringify({
             text: text, from_lang: fromLang, to_lang: toLang,