edit

app.py

@@ -11,35 +11,56 @@ from transformers import (
     TrOCRProcessor,
     VisionEncoderDecoderModel,
     pipeline,
-    AutoTokenizer,
+    VitsTokenizer,
     VitsModel,
 )
 
+# ============================
+# 1. Настройки устройства
+# ============================
 
-device_string: str = "cpu"
+device_string: str = "cuda" if torch.cuda.is_available() else "cpu"
 
+
+# ============================
+# 2. Модели
+# ============================
+
+# OCR: печатный английский текст
+# Модель: microsoft/trocr-small-printed
 ocr_processor: TrOCRProcessor = TrOCRProcessor.from_pretrained(
-    "raxtemur/trocr-base-ru"
+    "microsoft/trocr-small-printed"
 )
 ocr_model: VisionEncoderDecoderModel = VisionEncoderDecoderModel.from_pretrained(
-    "raxtemur/trocr-base-ru"
+    "microsoft/trocr-small-printed"
 )
 ocr_model.to(device_string)
 
+# Суммаризация: английский новостной/общий текст
+# Модель: sshleifer/distilbart-cnn-12-6
 summary_pipeline = pipeline(
     task="summarization",
-    model="IlyaGusev/mbart_ru_sum_gazeta",
-    tokenizer="IlyaGusev/mbart_ru_sum_gazeta",
+    model="sshleifer/distilbart-cnn-12-6",
 )
 
-tts_model: VitsModel = VitsModel.from_pretrained("facebook/mms-tts-rus")
-tts_tokenizer: AutoTokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-rus")
+# TTS: английская MMS VITS
+# Модель: facebook/mms-tts-eng
+tts_model: VitsModel = VitsModel.from_pretrained("facebook/mms-tts-eng")
+tts_tokenizer: VitsTokenizer = VitsTokenizer.from_pretrained("facebook/mms-tts-eng")
 tts_model.to(device_string)
 
+
+# ============================
+# 3. OCR
+# ============================
+
 def run_ocr(image_object: Image.Image) -> str:
     """
-    Распознавание текста с изображения.
-    Используем русскую TrOCR-модель.
+    Распознавание печатного английского текста с изображения.
+    Используем TrOCR (microsoft/trocr-small-printed).
+
+    Ожидается более-менее читаемый printed text
+    (скриншоты, документы, слайды и т.п.).
     """
     if image_object is None:
         return ""
@@ -52,28 +73,38 @@ def run_ocr(image_object: Image.Image) -> str:
     )
     pixel_values_tensor = processor_output.pixel_values.to(device_string)
 
-    generated_id_tensor = ocr_model.generate(pixel_values_tensor)
+    with torch.no_grad():
+        generated_id_tensor = ocr_model.generate(pixel_values_tensor)
+
     decoded_text_list = ocr_processor.batch_decode(
         generated_id_tensor,
         skip_special_tokens=True,
     )
 
-    recognized_text: str = decoded_text_list[0]
-    return recognized_text.strip()
+    recognized_text: str = decoded_text_list[0].strip()
+    return recognized_text
+
+
+# ============================
+# 4. Суммаризация (английский)
+# ============================
 
 def run_summarization(
     input_text: str,
     max_summary_tokens: int = 128,
 ) -> str:
     """
-    Русская суммаризация.
-    Без разбиения на чанки, так что огромные тексты лучше не подавать.
+    Английская суммаризация.
+    Без разбиения на чанки, поэтому очень длинные тексты лучше не подавать.
     """
     cleaned_text: str = input_text.strip()
     if not cleaned_text:
         return ""
 
     word_count: int = len(cleaned_text.split())
+
+    # Простая адаптация длины под размер текста,
+    # чтобы не было бессмысленных max_length >> input_length.
     dynamic_max_length: int = min(
         max_summary_tokens,
         max(32, word_count + 20),
@@ -82,7 +113,7 @@ def run_summarization(
     summary_result_list = summary_pipeline(
         cleaned_text,
         max_length=dynamic_max_length,
-        min_length=max(16, dynamic_max_length // 3),
+        min_length=max(10, dynamic_max_length // 3),
         do_sample=False,
     )
 
@@ -90,13 +121,17 @@ def run_summarization(
     return summary_text
 
 
+# ============================
+# 5. TTS (английский, MMS VITS)
+# ============================
+
 def run_tts(summary_text: str) -> Optional[str]:
     """
-    Озвучка текста конспекта через VitsModel (facebook/mms-tts-rus).
+    Озвучка английского текста конспекта через VitsModel (facebook/mms-tts-eng).
 
     ВАЖНО:
-    - защищаемся от пустого/битого ввода;
-    - ловим RuntimeError изнутри модели (известные проблемы MMS VITS на некоторых входах);
+    - защищаемся от пустого ввода;
+    - ловим RuntimeError изнутри модели (бывают краши на редких входах);
       в это�� случае просто возвращаем None, чтобы не ронять весь Space.
     """
     cleaned_text: str = summary_text.strip()
@@ -107,7 +142,11 @@ def run_tts(summary_text: str) -> Optional[str]:
         cleaned_text,
         return_tensors="pt",
     )
-    tokenized_inputs = {key: value.to(device_string) for key, value in tokenized_inputs.items()}
+
+    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:
@@ -140,15 +179,20 @@ def run_tts(summary_text: str) -> Optional[str]:
 
     return file_path
 
+
+# ============================
+# 6. Полный пайплайн
+# ============================
+
 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 не смог)
+    1) OCR: изображение -> исходный текст (английский)
+    2) Суммаризация: текст -> краткое резюме
+    3) TTS: резюме -> .wav файл (или None, если TTS не смог)
     """
     recognized_text: str = run_ocr(image_object=image_object)
 
@@ -162,42 +206,44 @@ def full_flow(
     return recognized_text, summary_text, audio_file_path
 
 
+# ============================
+# 7. Gradio UI
+# ============================
+
 gradio_interface = gradio_module.Interface(
     fn=full_flow,
     inputs=[
         gradio_module.Image(
             type="pil",
-            label="Изображение с текстом (желательно русский/английский, печатный)",
+            label="Image with printed English text",
         ),
         gradio_module.Slider(
             minimum=32,
             maximum=256,
             value=128,
             step=16,
-            label="Максимальная длина конспекта (токены, примерно)",
+            label="Maximum summary length (tokens, approx)",
         ),
     ],
     outputs=[
         gradio_module.Textbox(
-            label="Распознанный текст (OCR)",
+            label="Recognized text (OCR)",
             lines=6,
         ),
         gradio_module.Textbox(
-            label="Конспект (суммаризация)",
+            label="Summary (English)",
             lines=6,
         ),
         gradio_module.Audio(
-            label="Озвучка конспекта (MMS VITS, ru)",
+            label="Summary narration (MMS VITS, en)",
             type="filepath",
         ),
     ],
-    title="Картинка → Текст → Конспект → Озвучка (русские модели)",
+    title="Image → Text → Summary → Speech (English models)",
     description=(
-        "1) Русский трансформер OCR распознаёт текст с картинки.\n"
-        "2) Русский трансформер суммаризации делает краткий пересказ.\n"
-        "3) VITS-модель MMS (facebook/mms-tts-rus) озвучивает конспект.\n\n"
-        "Если озвучка не сгенерировалась, значит конкретный текст не понравился TTS-модели "
-        "и она упала внутри — пайплайн просто пропустит аудио."
+        "1) English OCR transformer recognizes printed text from the image.\n"
+        "2) English summarization transformer creates a short summary.\n"
+        "3) English VITS (facebook/mms-tts-eng) reads the summary aloud."
     ),
 )