Sandiago21's picture
Upload folder using huggingface_hub
e44712b
import gradio as gr
import numpy as np
import torch
from datasets import load_dataset
from transformers import SpeechT5ForTextToSpeech, SpeechT5HifiGan, SpeechT5Processor, pipeline
device = "cuda:0" if torch.cuda.is_available() else "cpu"
# load speech translation checkpoint
asr_pipe = pipeline("automatic-speech-recognition", model="openai/whisper-large-v2", device=device)
greek_translation_pipe = pipeline("translation", model="Helsinki-NLP/opus-mt-en-el")
# load text-to-speech checkpoint and speaker embeddings
model_id = "microsoft/speecht5_tts" # update with your model id
# pipe = pipeline("automatic-speech-recognition", model=model_id)
model = SpeechT5ForTextToSpeech.from_pretrained(model_id)
vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")
embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
speaker_embeddings = torch.tensor(embeddings_dataset[7440]["xvector"]).unsqueeze(0)
processor = SpeechT5Processor.from_pretrained(model_id)
model_id_greek = "Sandiago21/speecht5_finetuned_google_fleurs_greek"
model_greek = SpeechT5ForTextToSpeech.from_pretrained(model_id_greek)
processor_greek = SpeechT5Processor.from_pretrained(model_id_greek)
replacements = [
("á", "a"),
("â", "a"),
("ã", "a"),
("í", "i"),
("á", "a"),
("í", "i"),
("ñ", "n"),
("ó", "o"),
("ú", "u"),
("ü", "u"),
("á", "a"),
("ç", "c"),
("è", "e"),
("ì", "i"),
("í", "i"),
("ò", "o"),
("ó", "o"),
("ù", "u"),
("ú", "u"),
("š", "s"),
("ï", "i"),
("à", "a"),
("â", "a"),
("ç", "c"),
("è", "e"),
("ë", "e"),
("î", "i"),
("ï", "i"),
("ô", "o"),
("ù", "u"),
("û", "u"),
("ü", "u"),
("ου", "u"),
("αυ", "af"),
("ευ", "ef"),
("ει", "i"),
("οι", "i"),
("αι", "e"),
("ού", "u"),
("εί", "i"),
("οί", "i"),
("αί", "e"),
("Ά", "A"),
("Έ", "E"),
("Ή", "H"),
("Ί", "I"),
("Ό", "O"),
("Ύ", "Y"),
("Ώ", "O"),
("ΐ", "i"),
("Α", "A"),
("Β", "B"),
("Γ", "G"),
("Δ", "L"),
("Ε", "Ε"),
("Ζ", "Z"),
("Η", "I"),
("Θ", "Th"),
("Ι", "I"),
("Κ", "K"),
("Λ", "L"),
("Μ", "M"),
("Ν", "N"),
("Ξ", "Ks"),
("Ο", "O"),
("Π", "P"),
("Ρ", "R"),
("Σ", "S"),
("Τ", "T"),
("Υ", "Y"),
("Φ", "F"),
("Χ", "X"),
("Ω", "O"),
("ά", "a"),
("έ", "e"),
("ή", "i"),
("ί", "i"),
("α", "a"),
("β", "v"),
("γ", "g"),
("δ", "d"),
("ε", "e"),
("ζ", "z"),
("η", "i"),
("θ", "th"),
("ι", "i"),
("κ", "k"),
("λ", "l"),
("μ", "m"),
("ν", "n"),
("ξ", "ks"),
("ο", "o"),
("π", "p"),
("ρ", "r"),
("ς", "s"),
("σ", "s"),
("τ", "t"),
("υ", "i"),
("φ", "f"),
("χ", "h"),
("ψ", "ps"),
("ω", "o"),
("ϊ", "i"),
("ϋ", "i"),
("ό", "o"),
("ύ", "i"),
("ώ", "o"),
("í", "i"),
("õ", "o"),
("Ε", "E"),
("Ψ", "Ps"),
]
def cleanup_text(text):
for src, dst in replacements:
text = text.replace(src, dst)
return text
def synthesize_speech(text):
text = cleanup_text(text)
inputs = processor(text=text, return_tensors="pt")
speech = model.generate_speech(inputs["input_ids"].to(device), speaker_embeddings.to(device), vocoder=vocoder)
return gr.Audio.update(value=(16000, speech.cpu().numpy()))
def translate_to_english(audio):
outputs = asr_pipe(audio, max_new_tokens=256, generate_kwargs={"task": "translate", "language": "english"})
return outputs["text"]
def synthesise_from_english(text):
text = cleanup_text(text)
inputs = processor(text=text, return_tensors="pt")
speech = model.generate_speech(inputs["input_ids"].to(device), speaker_embeddings.to(device), vocoder=vocoder)
return speech.cpu().numpy()
def translate_from_english_to_greek(text):
return greek_translation_pipe(text)[0]["translation_text"]
def synthesise_from_greek(text):
text = cleanup_text(text)
inputs = processor_greek(text=text, return_tensors="pt")
speech = model_greek.generate_speech(inputs["input_ids"].to(device), speaker_embeddings.to(device), vocoder=vocoder)
return speech.cpu()
def speech_to_speech_translation(audio):
translated_text = translate_to_english(audio)
translated_text = translate_from_english_to_greek(translated_text)
# synthesised_speech = synthesise_from_english(translated_text)
# translated_text = translate_from_english_to_greek(synthesised_speech)
synthesised_speech = synthesise_from_greek(translated_text)
synthesised_speech = (synthesised_speech.numpy() * 32767).astype(np.int16)
return ((16000, synthesised_speech), translated_text)
title = "Cascaded STST"
description = """
Demo for cascaded speech-to-speech translation (STST), mapping from source speech in any language to target speech in Greek. Demo uses OpenAI's [Whisper Large v2](https://huggingface.co/openai/whisper-large-v2) model for speech translation, and [Sandiago21/speecht5_finetuned_google_fleurs_greek](https://huggingface.co/Sandiago21/speecht5_finetuned_google_fleurs_greek) checkpoint for text-to-speech, which is based on Microsoft's
[SpeechT5 TTS](https://huggingface.co/microsoft/speecht5_tts) model for text-to-speech, fine-tuned in Greek Audio dataset:
![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"), gr.outputs.Textbox()],
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"), gr.outputs.Textbox()],
examples=[["./example.wav"]],
title=title,
description=description,
)
with demo:
gr.TabbedInterface([mic_translate, file_translate], ["Microphone", "Audio File"])
demo.launch()