File size: 4,317 Bytes
e61ddf0 |
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 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 |
import gradio as gr
import numpy as np
import torch
import torchaudio
from transformers import SeamlessM4Tv2Model, AutoProcessor
from lang_list import (
ASR_TARGET_LANGUAGE_NAMES,
LANGUAGE_NAME_TO_CODE,
S2ST_TARGET_LANGUAGE_NAMES,
S2TT_TARGET_LANGUAGE_NAMES,
T2ST_TARGET_LANGUAGE_NAMES,
T2TT_TARGET_LANGUAGE_NAMES,
TEXT_SOURCE_LANGUAGE_NAMES,
)
processor = AutoProcessor.from_pretrained("facebook/seamless-m4t-v2-large")
model = SeamlessM4Tv2Model.from_pretrained("facebook/seamless-m4t-v2-large")
AUDIO_SAMPLE_RATE = 16000.0
MAX_INPUT_AUDIO_LENGTH = 60 # in seconds
DEFAULT_TARGET_LANGUAGE = "French"
if torch.cuda.is_available():
device = torch.device("cuda:0")
dtype = torch.float16
else:
device = torch.device("cpu")
dtype = torch.float32
def preprocess_audio(input_audio: str) -> None:
arr, org_sr = torchaudio.load(input_audio)
new_arr = torchaudio.functional.resample(arr, orig_freq=org_sr, new_freq=AUDIO_SAMPLE_RATE)
max_length = int(MAX_INPUT_AUDIO_LENGTH * AUDIO_SAMPLE_RATE)
if new_arr.shape[1] > max_length:
new_arr = new_arr[:, :max_length]
gr.Warning(f"Input audio is too long. Only the first {MAX_INPUT_AUDIO_LENGTH} seconds is used.")
torchaudio.save(input_audio, new_arr, sample_rate=int(AUDIO_SAMPLE_RATE))
def run_s2st(
input_audio: str, source_language: str, target_language: str
) -> tuple[tuple[int, np.ndarray] | None, str]:
preprocess_audio(input_audio)
source_language_code = LANGUAGE_NAME_TO_CODE[source_language]
target_language_code = LANGUAGE_NAME_TO_CODE[target_language]
arr, org_sr = torchaudio.load(input_audio)
audio_inputs = processor(audios=arr, return_tensors="pt",
sampling_rate=model.config.sampling_rate).to(device)
output = model.generate(**audio_inputs, return_intermediate_token_ids=True,
tgt_lang=target_language_code,)
audio_array_from_audio = output[0].cpu().numpy().squeeze()
text_tokens = output[2]
translated_text_from_text = processor.decode(text_tokens.tolist()[0], skip_special_tokens=True)
return (int(AUDIO_SAMPLE_RATE), audio_array_from_audio), translated_text_from_text
description = """
# Direct Speech to Speech Translation
This demo uses SeamlessM4T V2 to translate one speech directly into another.
The model being used here is [facebook/seamless-m4t-v2-large](https://huggingface.co/facebook/seamless-m4t-v2-large).
SeamlessM4T V2 is unified model enables multiple tasks like Speech-to-Speech (S2ST), Speech-to-Text (S2TT), Text-to-Speech (T2ST) translation and more, without relying on multiple separate models.
"""
with gr.Blocks() as demo_s2st:
gr.Markdown(description)
with gr.Row():
with gr.Column():
with gr.Group():
input_audio = gr.Audio(label="Input speech", type="filepath")
source_language = gr.Dropdown(
label="Source language",
choices=ASR_TARGET_LANGUAGE_NAMES,
value="English",
)
target_language = gr.Dropdown(
label="Target language",
choices=S2ST_TARGET_LANGUAGE_NAMES,
value=DEFAULT_TARGET_LANGUAGE,
)
btn = gr.Button("Translate")
with gr.Column():
with gr.Group():
output_audio = gr.Audio(
label="Translated speech",
autoplay=False,
streaming=False,
type="numpy",
)
output_text = gr.Textbox(label="Translated text")
gr.Examples(
examples=[
["assets/sample_input.mp3", "English", "French"],
["assets/sample_input.mp3", "English", "Mandarin Chinese"],
["assets/sample_input_2.mp3", "English", "Hindi"],
["assets/sample_input_2.mp3", "English", "Spanish"],
],
inputs=[input_audio, source_language, target_language],
outputs=[output_audio, output_text],
fn=run_s2st,
cache_examples=True,
)
btn.click(
fn=run_s2st,
inputs=[input_audio, source_language, target_language],
outputs=[output_audio, output_text],
)
demo_s2st.launch() |