File size: 6,160 Bytes
291564a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e796a00
291564a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d0aa8cd
e796a00
 
291564a
 
 
 
 
 
 
1180d04
 
 
 
 
 
 
 
 
 
 
 
291564a
d0aa8cd
291564a
 
1180d04
291564a
99ac92b
 
 
 
 
 
 
 
 
 
5fad5d4
99ac92b
1180d04
 
 
 
 
291564a
1180d04
291564a
 
 
 
 
 
1180d04
e796a00
291564a
e796a00
d0aa8cd
e796a00
291564a
 
 
 
 
 
 
 
0487019
 
 
 
 
5fad5d4
 
0487019
 
306b289
 
 
79b79bf
306b289
 
 
79b79bf
306b289
 
291564a
 
219e01a
291564a
219e01a
 
1428e8e
219e01a
 
 
291564a
 
 
 
 
 
1180d04
291564a
3c2c1d3
 
 
 
 
 
 
 
 
 
 
 
af003e9
1180d04
48f97c1
 
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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
import gradio as gr
import torch
import uuid
import json
import librosa
import os
import tempfile
import soundfile as sf
import scipy.io.wavfile as wav

from transformers import pipeline, VitsModel, AutoTokenizer, set_seed
from nemo.collections.asr.models import EncDecMultiTaskModel

# Constants
SAMPLE_RATE = 16000  # Hz

# load ASR model
canary_model = EncDecMultiTaskModel.from_pretrained('nvidia/canary-1b')
decode_cfg = canary_model.cfg.decoding
decode_cfg.beam.beam_size = 1
canary_model.change_decoding_strategy(decode_cfg)

# Function to convert audio to text using ASR
def gen_text(audio_filepath, action, source_lang, target_lang):
    if audio_filepath is None:
        raise gr.Error("Please provide some input audio.")
    
    utt_id = uuid.uuid4()
    with tempfile.TemporaryDirectory() as tmpdir:
        # Convert to 16 kHz
        data, sr = librosa.load(audio_filepath, sr=None, mono=True)
        if sr != SAMPLE_RATE:
            data = librosa.resample(data, orig_sr=sr, target_sr=SAMPLE_RATE)
        converted_audio_filepath = os.path.join(tmpdir, f"{utt_id}.wav")
        sf.write(converted_audio_filepath, data, SAMPLE_RATE)

        # Transcribe audio
        duration = len(data) / SAMPLE_RATE
        manifest_data = {
            "audio_filepath": converted_audio_filepath,
            "taskname": action,
            "source_lang": source_lang,
            "target_lang": source_lang if action=="asr" else target_lang,
            "pnc": "no",
            "answer": "predict",
            "duration": str(duration),
        }
        manifest_filepath = os.path.join(tmpdir, f"{utt_id}.json")
        with open(manifest_filepath, 'w') as fout:
            fout.write(json.dumps(manifest_data))

        predicted_text = canary_model.transcribe(manifest_filepath)[0]
        # if duration < 40:
        #     predicted_text = canary_model.transcribe(manifest_filepath)[0]
        # else:
        #     predicted_text = get_buffered_pred_feat_multitaskAED(
        #         frame_asr,
        #         canary_model.cfg.preprocessor,
        #         model_stride_in_secs,
        #         canary_model.device,
        #         manifest=manifest_filepath,
        #     )[0].text
    
    return predicted_text

# Function to convert text to speech using TTS
def gen_speech(text, lang):
    set_seed(555)  # Make it deterministic
    match lang:
        case "en":
             model = "facebook/mms-tts-eng"
        case "fr":
             model = "facebook/mms-tts-fra"
        case "de":
             model = "facebook/mms-tts-deu"
        case "es":
             model = "facebook/mms-tts-spa"
        case _:
            model = "facebook/mms-tts"
    
    # load TTS model
    tts_model = VitsModel.from_pretrained(model)
    tts_tokenizer = AutoTokenizer.from_pretrained(model)

    input_text = tts_tokenizer(text, return_tensors="pt")
    with torch.no_grad():
        outputs = tts_model(**input_text)
    waveform_np = outputs.waveform[0].cpu().numpy()
    output_file = f"{str(uuid.uuid4())}.wav"
    wav.write(output_file, rate=tts_model.config.sampling_rate, data=waveform_np)
    return output_file

# Root function for Gradio interface
def start_process(audio_filepath, source_lang, target_lang):
    transcription = gen_text(audio_filepath, "asr", source_lang, target_lang)
    print("Done transcribing")
    translation = gen_text(audio_filepath, "s2t_translation", source_lang, target_lang) 
    print("Done translation")
    audio_output_filepath = gen_speech(translation, target_lang)
    print("Done speaking")   
    return transcription, translation, audio_output_filepath

    
# Create Gradio interface
playground = gr.Blocks()

with playground:

    with gr.Row():
        gr.Markdown("""
                    ## Your AI Translate Assistant
                    ### Gets input audio from user, transcribe and translate it. Convert back to speech.
                    - category: [Automatic Speech Recognition](https://huggingface.co/models?pipeline_tag=automatic-speech-recognition), model: [nvidia/canary-1b](https://huggingface.co/nvidia/canary-1b)
                    - category: [Text-to-Speech](https://huggingface.co/models?pipeline_tag=text-to-speech), model: [facebook/mms-tts](https://huggingface.co/facebook/mms-tts)
                    """)
        
    with gr.Row():
        with gr.Column():
            source_lang = gr.Dropdown(
                choices=["en", "de", "es", "fr"], value="en", label="Source Language"
            )
        with gr.Column():
            target_lang = gr.Dropdown(
                choices=["en", "de", "es", "fr"], value="fr", label="Target Language"
            )            

    with gr.Row():
        with gr.Column():
            input_audio = gr.Audio(sources=["microphone"], type="filepath", label="Input Audio")            
        with gr.Column():
            translated_speech = gr.Audio(type="filepath", label="Generated Speech")            

    with gr.Row():
        with gr.Column():            
            transcipted_text = gr.Textbox(label="Transcription")
        with gr.Column():            
            translated_text = gr.Textbox(label="Translation")

    with gr.Row():
        with gr.Column():
            submit_button = gr.Button(value="Start Process", variant="primary")
        with gr.Column():
            clear_button = gr.ClearButton(components=[input_audio, source_lang, target_lang, transcipted_text, translated_text, translated_speech], value="Clear")    

    with gr.Row():
        gr.Examples(
            examples=[
                ["sample_en.wav","en","fr"],
                ["sample_fr.wav","fr","de"],
                ["sample_de.wav","de","es"],
                ["sample_es.wav","es","en"]
            ], 
            inputs=[input_audio, source_lang, target_lang], 
            outputs=[transcipted_text, translated_text, translated_speech], 
            run_on_click=True, cache_examples=True, fn=start_process
        )
            
    submit_button.click(start_process, inputs=[input_audio, source_lang, target_lang], outputs=[transcipted_text, translated_text, translated_speech])

playground.launch()