File size: 3,101 Bytes
d347764 ff3f816 d347764 9e2a87c 0c8ad01 d347764 0c8ad01 9e2a87c d347764 e970d56 472a003 0c8ad01 d347764 9e2a87c 6ab6711 d347764 0c8ad01 bc20e42 e467863 be03dff d347764 6ab6711 d614113 af4a9d1 d347764 f805e49 d386a01 f805e49 c737803 d347764 226ec3a d347764 f805e49 d347764 c737803 3946ba6 c737803 49c1298 9e2a87c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
import gradio as gr
import logging
import numpy as np
import torch
from transformers import VitsModel, VitsTokenizer, pipeline
from transformers import M2M100ForConditionalGeneration
from tokenization_small100 import SMALL100Tokenizer
device = "cuda:0" if torch.cuda.is_available() else "cpu"
target_language = "fr"
# load speech translation checkpoint
asr_pipe = pipeline("automatic-speech-recognition", model="bofenghuang/whisper-small-cv11-french", device=device)
translation_model = M2M100ForConditionalGeneration.from_pretrained("alirezamsh/small100")
translation_tokenizer = SMALL100Tokenizer.from_pretrained("alirezamsh/small100", tgt_lang=target_language)
# load text-to-speech checkpoint
model = VitsModel.from_pretrained("facebook/mms-tts-fra")
tokenizer = VitsTokenizer.from_pretrained("facebook/mms-tts-fra")
def translate(audio):
outputs = asr_pipe(audio, max_new_tokens=256, generate_kwargs={"task": "translate"})
eng_text = outputs["text"]
encoded_eng_text = translation_tokenizer(eng_text, return_tensors="pt")
generated_tokens = translation_model.generate(**encoded_eng_text)
translated_text = translation_tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
logging.info(f"Translated Text: {translated_text}")
return translated_text
def synthesise(text):
inputs = tokenizer(text, return_tensors="pt")
with torch.no_grad():
outputs = model(inputs["input_ids"])
speech = outputs["waveform"][0]
logging.info(speech)
return speech.cpu()
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
title = "Cascaded STST"
description = """
Demo for cascaded speech-to-speech translation (STST), mapping from source speech in any language to target speech in French. Demo uses OpenAI's [Whisper Base](https://huggingface.co/openai/whisper-base) model for ASR, the
[SMaLL-100](https://huggingface.co/alirezamsh/small100) model for text to text translation and Facebook's[MMS TTS-FRA](https://huggingface.co/facebook/mms-tts-fra) for text-to-speech for french:
"""
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=speech_to_speech_translation,
inputs=gr.Audio(source="upload", type="filepath"),
outputs=gr.Audio(label="Generated Speech", type="numpy"),
examples=[["./example.wav"]],
title=title,
description=description,
)
with demo:
gr.TabbedInterface([mic_translate, file_translate], ["Microphone", "Audio File"])
logging.getLogger().setLevel(logging.INFO)
demo.launch() |