Update app.py

app.py

@@ -4,8 +4,8 @@ from PIL import Image
 from transformers import (
     BlipProcessor,
     BlipForConditionalGeneration,
-    AutoTokenizer,
-    AutoModelForSeq2SeqLM
+    M2M100Tokenizer,
+    M2M100ForConditionalGeneration
 )
 from typing import Union
 from gtts import gTTS
@@ -14,8 +14,10 @@ import uuid
 import time
 import gc
 
+# Оптимизация CPU: установка числа потоков
 torch.set_num_threads(2)
 
+# Глобальная переменная для кэширования pipeline
 _pipeline = None
 
 def init_pipeline():
@@ -33,8 +35,8 @@ class ImageCaptionPipeline:
         print(f"Время загрузки BLIP: {time.time() - start_time:.2f} секунд")
         
         start_time = time.time()
-        self.translator_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-ru")
-        self.translator_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-ru").to(self.device)
+        self.translator_tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_418M")
+        self.translator_model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_418M").to(self.device)
         print(f"Время загрузки переводчика: {time.time() - start_time:.2f} секунд")
 
     def generate_captions(self, image: Union[str, Image.Image]) -> tuple:
@@ -42,6 +44,7 @@ class ImageCaptionPipeline:
         if isinstance(image, str):
             image = Image.open(image)
         image = image.convert("RGB")
+        image = image.resize((384, 384))  # Рекомендованный размер для BLIP-large
         inputs = self.blip_processor(images=image, return_tensors="pt").to(self.device)
         with torch.no_grad():
             output_ids = self.blip_model.generate(**inputs, max_length=50, num_beams=2, early_stopping=True)
@@ -49,9 +52,16 @@ class ImageCaptionPipeline:
         print(f"Время генерации английской подписи: {time.time() - start_time:.2f} секунд")
         
         start_time = time.time()
+        self.translator_tokenizer.src_lang = "en"
         translated_inputs = self.translator_tokenizer(english_caption, return_tensors="pt", padding=True).to(self.device)
         with torch.no_grad():
-            translated_ids = self.translator_model.generate(**translated_inputs, max_length=50, num_beams=2, early_stopping=True)
+            translated_ids = self.translator_model.generate(
+                **translated_inputs,
+                forced_bos_token_id=self.translator_tokenizer.get_lang_id("ru"),
+                max_length=50,
+                num_beams=2,
+                early_stopping=True
+            )
             russian_caption = self.translator_tokenizer.decode(translated_ids[0], skip_special_tokens=True)
         print(f"Время перевода на русский: {time.time() - start_time:.2f} секунд")
         
@@ -79,21 +89,21 @@ def generate_audio(english_caption: str, russian_caption: str, audio_language: s
     if not english_caption and not russian_caption:
         return None
     pipeline = init_pipeline()
-    text = russian_caption.replace("Русский: ", "") if audio_language == "Русский" else english_caption.replace("English: ", "")
+    text = russian_caption if audio_language == "Русский" else english_caption
     return pipeline.generate_audio(text, audio_language)
 
-with gr.Blocks() as iface:
+with gr.Blocks(css=".btn {width: 200px; background-color: #4682B4; color: white; border: none; padding: 10px 20px; text-align: center; font-size: 16px;} .equal-height { height: 40px; }") as iface:
     with gr.Row():
-        with gr.Column(scale=1, min_width=600, variant="panel"):
-            image = gr.Image(type="pil", label="Изображение", height=600, width=600)
-            submit_button = gr.Button("Сгенерировать описание", elem_classes="btn", size="sm")
+        with gr.Column(scale=1, min_width=400, variant="panel"):
+            image = gr.Image(type="pil", label="Изображение", height=400, width=400)
+            submit_button = gr.Button("Сгенерировать описание", elem_classes="btn")
         with gr.Column(scale=1, min_width=300):
             english_caption = gr.Textbox(label="Подпись English:", lines=2)
             russian_caption = gr.Textbox(label="Подпись Русский:", lines=2)
-            audio_button = gr.Button("Сгенерировать озвучку", elem_classes="btn", size="sm")
+            audio_button = gr.Button("Сгенерировать озвучку", elem_classes="btn")
             with gr.Row():
-                audio_language = gr.Dropdown(choices=["Русский", "English"], label="Язык озвучки", value="Русский", scale=1)
-                audio_output = gr.Audio(label="Озвучка", scale=1, min_width=150)
+                audio_language = gr.Dropdown(choices=["Русский", "English"], label="Язык озвучки", value="Русский", scale=1, min_width=150, elem_classes="equal-height")
+                audio_output = gr.Audio(label="Озвучка", scale=1, min_width=150, elem_classes="equal-height")
     
     submit_button.click(
         fn=generate_captions,