edit

app.py

@@ -1,9 +1,11 @@
 from typing import Tuple, Optional
+
 import tempfile
-import soundfile as sf
+
+import numpy as numpy_module
+import soundfile as soundfile_module
 import torch
-import gradio as gr
-import numpy as np
+import gradio as gradio_module
 from PIL import Image
 from transformers import (
     TrOCRProcessor,
@@ -14,28 +16,30 @@ from transformers import (
 )
 
 
+device_string: str = "cpu"
+
 ocr_processor: TrOCRProcessor = TrOCRProcessor.from_pretrained(
     "raxtemur/trocr-base-ru"
 )
 ocr_model: VisionEncoderDecoderModel = VisionEncoderDecoderModel.from_pretrained(
     "raxtemur/trocr-base-ru"
 )
-ocr_model.to("cpu")
+ocr_model.to(device_string)
 
 summary_pipeline = pipeline(
-    "summarization",
+    task="summarization",
     model="IlyaGusev/mbart_ru_sum_gazeta",
     tokenizer="IlyaGusev/mbart_ru_sum_gazeta",
 )
 
-tts_model = VitsModel.from_pretrained("facebook/mms-tts-rus")
-tts_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-rus")
-tts_model.to("cpu")
+tts_model: VitsModel = VitsModel.from_pretrained("facebook/mms-tts-rus")
+tts_tokenizer: AutoTokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-rus")
+tts_model.to(device_string)
 
 def run_ocr(image_object: Image.Image) -> str:
     """
     Распознавание текста с изображения.
-    Предполагаем, что на картинке русский/кириллический или латинский печатный текст.
+    Используем русскую TrOCR-модель.
     """
     if image_object is None:
         return ""
@@ -46,7 +50,7 @@ def run_ocr(image_object: Image.Image) -> str:
         images=rgb_image_object,
         return_tensors="pt",
     )
-    pixel_values_tensor = processor_output.pixel_values.to("cpu")
+    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(
@@ -57,48 +61,146 @@ def run_ocr(image_object: Image.Image) -> str:
     recognized_text: str = decoded_text_list[0]
     return recognized_text.strip()
 
-def run_summary(text: str) -> str:
-    text = text.strip()
-    if not text:
+def run_summarization(
+    input_text: str,
+    max_summary_tokens: int = 128,
+) -> str:
+    """
+    Русская суммаризация.
+    Без разбиения на чанки, так что огромные тексты лучше не подавать.
+    """
+    cleaned_text: str = input_text.strip()
+    if not cleaned_text:
         return ""
-    result = summary_pipeline(text, max_length=128, min_length=30, do_sample=False)
-    return result[0]["summary_text"].strip()
 
-def run_tts(text: str) -> Optional[str]:
-    text = text.strip()
-    if not text:
+    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=dynamic_max_length,
+        min_length=max(16, dynamic_max_length // 3),
+        do_sample=False,
+    )
+
+    summary_text: str = summary_result_list[0]["summary_text"].strip()
+    return summary_text
+
+
+def run_tts(summary_text: str) -> Optional[str]:
+    """
+    Озвучка текста конспекта через VitsModel (facebook/mms-tts-rus).
+
+    ВАЖНО:
+    - защищаемся от пустого/битого ввода;
+    - ловим RuntimeError изнутри модели (известные проблемы MMS VITS на некоторых входах);
+      в этом случае просто возвращаем None, чтобы не ронять весь Space.
+    """
+    cleaned_text: str = summary_text.strip()
+    if not cleaned_text:
+        return None
+
+    tokenized_inputs = tts_tokenizer(
+        cleaned_text,
+        return_tensors="pt",
+    )
+    tokenized_inputs = {key: value.to(device_string) for key, value in tokenized_inputs.items()}
+
+    input_ids_tensor = tokenized_inputs.get("input_ids")
+    if input_ids_tensor is None:
         return None
-    inputs = tts_tokenizer(text, return_tensors="pt").to("cpu")
-    with torch.no_grad():
-        waveform = tts_model(**inputs).waveform
-    audio = waveform.squeeze().cpu().numpy().astype("float32")
-    audio = np.clip(audio, -1.0, 1.0)
-    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
-        sf.write(f.name, audio, tts_model.config.sampling_rate)
-        return f.name
-
-def full_flow(image: Image.Image) -> Tuple[str, str, Optional[str]]:
-    text = run_ocr(image)
-    summary = run_summary(text)
-    audio_path = run_tts(summary)
-    return text, summary, audio_path
-
-
-demo = gr.Interface(
+
+    if input_ids_tensor.numel() == 0 or input_ids_tensor.shape[1] == 0:
+        return None
+
+    try:
+        with torch.no_grad():
+            model_output = tts_model(**tokenized_inputs)
+            waveform_tensor = model_output.waveform  # shape: (batch, n_samples)
+    except RuntimeError as runtime_error:
+        print(f"[WARN] TTS RuntimeError: {runtime_error}")
+        return None
+
+    waveform_array = waveform_tensor.squeeze().cpu().numpy().astype("float32")
+    waveform_array = numpy_module.clip(waveform_array, -1.0, 1.0)
+
+    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[str]]:
+    """
+    Полный пайплайн:
+    1) OCR: изображение -> исходный текст
+    2) Суммаризация: текст -> конспект
+    3) TTS: конспект -> .wav файл (или None, если TTS не смог)
+    """
+    recognized_text: str = run_ocr(image_object=image_object)
+
+    summary_text: str = run_summarization(
+        input_text=recognized_text,
+        max_summary_tokens=max_summary_tokens,
+    )
+
+    audio_file_path: Optional[str] = run_tts(summary_text=summary_text)
+
+    return recognized_text, summary_text, audio_file_path
+
+
+gradio_interface = gradio_module.Interface(
     fn=full_flow,
-    inputs=gr.Image(type="pil", label="Изображение с текстом (русский или английский)"),
+    inputs=[
+        gradio_module.Image(
+            type="pil",
+            label="Изображение с текстом (желательно русский/английский, печатный)",
+        ),
+        gradio_module.Slider(
+            minimum=32,
+            maximum=256,
+            value=128,
+            step=16,
+            label="Максимальная длина конспекта (токены, примерно)",
+        ),
+    ],
     outputs=[
-        gr.Textbox(label="Распознанный текст", lines=6),
-        gr.Textbox(label="Краткий пересказ", lines=6),
-        gr.Audio(label="Озвучка конспекта", type="filepath"),
+        gradio_module.Textbox(
+            label="Распознанный текст (OCR)",
+            lines=6,
+        ),
+        gradio_module.Textbox(
+            label="Конспект (суммаризация)",
+            lines=6,
+        ),
+        gradio_module.Audio(
+            label="Озвучка конспекта (MMS VITS, ru)",
+            type="filepath",
+        ),
     ],
-    title="Картинка → Текст → Конспект → Озвучка (русская версия)",
+    title="Картинка → Текст → Конспект → Озвучка (русские модели)",
     description=(
-        "1️⃣ OCR (TrOCR-base) распознаёт текст с картинки.\n"
-        "2️⃣ Суммаризация (IlyaGusev/mbart_ru_sum_gazeta) делает конспект.\n"
-        "3️⃣ TTS (facebook/mms-tts-rus) озвучивает результат."
+        "1) Русский трансформер OCR распознаёт текст с картинки.\n"
+        "2) Русский трансформер суммаризации делает краткий пересказ.\n"
+        "3) VITS-модель MMS (facebook/mms-tts-rus) озвучивает конспект.\n\n"
+        "Если озвучка не сгенерировалась, значит конкретный текст не понравился TTS-модели "
+        "и она упала внутри — пайплайн просто пропустит аудио."
     ),
 )
 
+
 if __name__ == "__main__":
-    demo.launch()
+    gradio_interface.launch()