Update app.py

app.py

@@ -3,66 +3,63 @@ import numpy as np
 import torch
 from datasets import load_dataset
 
-from transformers import AutoTokenizer, VitsModel, pipeline as huggingface_pipeline
+from transformers import AutoTokenizer, VitsModel, pipeline
 
-# Проверяем доступность CUDA для ускорения вычислений
-device_name = "cuda:0" if torch.cuda.is_available() else "cpu"
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
-# Загружаем модели для распознавания речи и перевода
-speech_recognition_model = huggingface_pipeline("automatic-speech-recognition", model="openai/whisper-tiny", device=device_name)
-my_awesome_translator = huggingface_pipeline("translation", model="Helsinki-NLP/opus-mt-en-ru")
+asr_pipe = pipeline("automatic-speech-recognition", model="openai/whisper-tiny", device=device) #Тут добавил tiny, потому что модель станет более компактной
+translater = pipeline("translation", model="Helsinki-NLP/opus-mt-en-ru")  # Инициализация модели для перевода текста на русский язык
 
-# Функция для перевода речи на русский язык
-def translate_audio(audio, translator_model=my_awesome_translator):
-    speech_recognition_result = speech_recognition_model(audio, max_new_tokens=256, generate_kwargs={"task": "translate"})
-    return translator_model(speech_recognition_result['text'])[0]['translation_text']
+def translate(audio, translater: pipeline = translater): # Определение функции для перевода аудио в текст
+    outputs = asr_pipe(audio, max_new_tokens=256, generate_kwargs={"task": "translate"}) # Получение текстового представления аудио
+    return translater(outputs['text'])[0]['translation_text'] # Возврат переведенного текста
 
-# Модели для синтеза голоса
-speech_synthesis_model = VitsModel.from_pretrained("facebook/mms-tts-rus")
-text_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-rus")
 
-# Функция для синтеза речи
-def custom_speech_synthesis(text, tokenizer=text_tokenizer, model=speech_synthesis_model):
-    input_tokens = tokenizer(text, return_tensors="pt")
+model = VitsModel.from_pretrained("facebook/mms-tts-rus") # Загрузка модели для генерации речи на русском языке
+tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-rus") # Загрузка токенизатора для модели
+
+def synthesise(text: str, tokenizer: AutoTokenizer = tokenizer, model: VitsModel = model): # Определение функции для синтеза речи из текста
+    inputs = tokenizer(text, return_tensors="pt") # Создание токенизированного представления текста
+    # print(inputs)
     with torch.no_grad():
-        synthesized_output = model(input_tokens).waveform
-    return synthesized_output.cpu()
-
-# Функция для обработки речи и ее перевода
-def process_speech_to_speech_translation(input_audio):
-    translated_text = translate_audio(input_audio)
-    synthesized_speech = custom_speech_synthesis(translated_text)
-    synthesized_speech = (synthesized_speech.numpy() * 32767).astype(np.int16)
-    return 16000, synthesized_speech[0]
-
-# Новое название и описание для интерфейса
-interface_title = "Speech Translation and Synthesis"
-interface_description = """
-Experience the magic of speech-to-speech translation! Our innovative system translates your speech and synthesizes it in Russian. This demo utilizes cutting-edge models for speech recognition, translation, and text-to-speech synthesis.
-"""
+        output = model(**inputs).waveform # Генерация аудиофайла из текста
+    return output.cpu()  # Возврат полученной речи
 
 
-demo = gr.Blocks()
+def speech_to_speech_translation(audio): # Определение функции для перевода и синтеза речи
+    translated_text = translate(audio) # Перевод аудио в текст
+    synthesised_speech = synthesise(translated_text) # Генерация речи на русском языке из переведенного текста
+    synthesised_speech = (synthesised_speech.numpy() * 32767).astype(np.int16) # Преобразование и нормализация речи
+    return 16000, synthesised_speech[0] # Возврат частоты дискретизации и синтезированной речи
 
-mic_translate = gr.Interface(
-    fn=process_speech_to_speech_translation,
-    inputs=gr.Audio(source="microphone", type="filepath"),
-    outputs=gr.Audio(label="Generated Speech", type="numpy"),
+
+title = "Cascaded STST"
+description = """
+Demo for cascaded speech-to-speech translation (STST), mapping from source speech in multi language to target speech in Russian. Demo uses OpenAI's [Whisper Tiny](https://huggingface.co/openai/whisper-tiny) model for speech translation, and Facebook's
+[mms-tts-rus](https://huggingface.co/acebook/mms-tts-rus) model for text-to-speech:
+![Cascaded STST](https://huggingface.co/datasets/huggingface-course/audio-course-images/resolve/main/s2st_cascaded.png "Diagram of cascaded speech to speech translation")
+"""
+
+demo = gr.Blocks() # Создание блока для интерфейса
+
+mic_translate = gr.Interface( # Создание интерфейса для микрофона
+    fn=speech_to_speech_translation,  # Используемая функция
+    inputs=gr.Audio(source="microphone", type="filepath"), # Ввод с микрофона
+    outputs=gr.Audio(label="Generated Speech", type="numpy"), # Вывод сгенерированной речи
     title=title,
     description=description,
 )
 
-# Создаем интерфейс
-file_translate = gr.Interface(
-    fn=process_speech_to_speech_translation,
-    inputs=gr.Audio(source="microphone", type="file"),
-    outputs=gr.Audio(label="Synthesized Speech", type="numpy"),
-    title=interface_title,
-    description=interface_description,
+file_translate = gr.Interface( # Создание интерфейса для загрузки аудиофайла
+    fn=speech_to_speech_translation,
+    inputs=gr.Audio(source="upload", type="filepath"),
+    outputs=gr.Audio(label="Generated Speech", type="numpy"),
+    examples=[["./test_2.wav"]],
+    title=title,
+    description=description,
 )
 
-# Запускаем интерфейс
 with demo:
     gr.TabbedInterface([mic_translate, file_translate], ["Microphone", "Audio File"])
 
-demo.launch()
+demo.launch()