edit

app.py

@@ -1,11 +1,18 @@
-from typing import Tuple, Optional, Any
+from typing import Tuple, Optional
+
+import tempfile
+
 import numpy as np
+import soundfile as soundfile_module
+import torch
 import gradio as gradio_module
 from PIL import Image
 from transformers import (
     TrOCRProcessor,
     VisionEncoderDecoderModel,
     pipeline,
+    VitsModel,
+    AutoTokenizer,
 )
 
 ocr_processor: TrOCRProcessor = TrOCRProcessor.from_pretrained(
@@ -20,10 +27,12 @@ summary_pipeline = pipeline(
     model="sshleifer/distilbart-cnn-12-6",
 )
 
-text_to_speech_pipeline = pipeline(
-    task="text-to-speech",
-    model="facebook/mms-tts-eng",
-)
+tts_model: VitsModel = VitsModel.from_pretrained("facebook/mms-tts-rus")
+tts_tokenizer: AutoTokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-rus")
+
+device_string: str = "cpu"
+ocr_model.to(device_string)
+tts_model.to(device_string)
 
 
 def run_ocr(image_object: Image.Image) -> str:
@@ -38,7 +47,7 @@ def run_ocr(image_object: Image.Image) -> str:
         images=image_object,
         return_tensors="pt",
     )
-    pixel_values_tensor = processor_output.pixel_values
+    pixel_values_tensor = processor_output.pixel_values.to(device_string)
 
     generated_id_tensor = ocr_model.generate(pixel_values_tensor)
     decoded_text_list = ocr_processor.batch_decode(
@@ -56,16 +65,22 @@ def run_summarization(
 ) -> str:
     """
     Суммаризация текста до короткого конспекта.
-    Без разбиения на чанки, поэтому огромные тексты лучше не подавать.
+    Без сложного разбиения на чанки -> длинные тексты лучше не кормить.
     """
     cleaned_text: str = input_text.strip()
     if not cleaned_text:
         return ""
 
+    word_count: int = len(cleaned_text.split())
+    dynamic_max_length: int = min(
+        max_summary_tokens,
+        max(32, word_count + 20),
+    )
+
     summary_result_list = summary_pipeline(
         cleaned_text,
-        max_length=max_summary_tokens,
-        min_length=max(16, max_summary_tokens // 3),
+        max_length=dynamic_max_length,
+        min_length=max(10, dynamic_max_length // 3),
         do_sample=False,
     )
 
@@ -73,30 +88,44 @@ def run_summarization(
     return summary_text
 
 
-def run_tts(summary_text: str) -> Optional[Tuple[int, Any]]:
+def run_tts(summary_text: str) -> Optional[str]:
     """
-    Озвучка текста конспекта.
-    Используем модель, которой не нужны внешние speaker embeddings.
+    Озвучка текста конспекта через VitsModel (facebook/mms-tts-rus).
+    Возвращаем путь до временного .wav файла, который Gradio отдаст в плеер.
     """
     cleaned_text: str = summary_text.strip()
     if not cleaned_text:
         return None
 
-    tts_output = text_to_speech_pipeline(cleaned_text)
+    tokenized_inputs = tts_tokenizer(
+        cleaned_text,
+        return_tensors="pt",
+    ).to(device_string)
 
-    sampling_rate_int: int = int(tts_output["sampling_rate"])
-    audio_array = tts_output["audio"]
+    with torch.no_grad():
+        model_output = tts_model(**tokenized_inputs)
+        waveform_tensor = model_output.waveform
 
-    audio_array = np.array(audio_array, dtype=np.float32)
-    audio_array = np.clip(audio_array, -1.0, 1.0)
+    waveform_array = waveform_tensor.squeeze().cpu().numpy().astype("float32")
 
-    return sampling_rate_int, audio_array
+    with tempfile.NamedTemporaryFile(
+        suffix=".wav",
+        delete=False,
+    ) as temporary_file:
+        soundfile_module.write(
+            temporary_file.name,
+            waveform_array,
+            tts_model.config.sampling_rate,
+        )
+        file_path: str = temporary_file.name
+
+    return file_path
 
 
 def full_flow(
     image_object: Image.Image,
     max_summary_tokens: int = 128,
-) -> Tuple[str, str, Optional[Tuple[int, Any]]]:
+) -> Tuple[str, str, Optional[str]]:
 
     recognized_text: str = run_ocr(image_object=image_object)
 
@@ -105,16 +134,17 @@ def full_flow(
         max_summary_tokens=max_summary_tokens,
     )
 
-    audio_tuple = run_tts(summary_text=summary_text)
+    audio_file_path: Optional[str] = run_tts(summary_text=summary_text)
+
+    return recognized_text, summary_text, audio_file_path
 
-    return recognized_text, summary_text, audio_tuple
 
 gradio_interface = gradio_module.Interface(
     fn=full_flow,
     inputs=[
         gradio_module.Image(
             type="pil",
-            label="Изображение с напечатанным текстом (английский)",
+            label="Изображение с напечатанным текстом (лучше русским/латиницей)",
         ),
         gradio_module.Slider(
             minimum=32,
@@ -134,15 +164,15 @@ gradio_interface = gradio_module.Interface(
             lines=6,
         ),
         gradio_module.Audio(
-            label="Озвучка конспекта (TTS)",
-            type="numpy",
+            label="Озвучка конспекта (VITS, ru)",
+            type="filepath",
         ),
     ],
     title="Картинка → Конспект → Озвучка (Transformers)",
     description=(
         "1) Трансформер OCR распознаёт текст с изображения. "
         "2) Трансформер суммаризации сокращает текст до конспекта. "
-        "3) Трансформер TTS озвучивает конспект."
+        "3) VITS-модель (facebook/mms-tts-rus) озвучивает конспект по-русски."
     ),
 )
 

requirements.txt

@@ -1,6 +1,7 @@
-transformers
+transformers>=4.33.0
 torch
 sentencepiece
 gradio
 Pillow
 numpy
+soundfile