from transformers import pipeline
import os
import gradio as gr
import torch

#Text to text
#translator = pipeline(task="translation",
#                      model="facebook/nllb-200-distilled-600M",
#                      torch_dtype=torch.bfloat16) 
#Text to audio
pipe = pipeline("text-to-speech", model="suno/bark-small")

                
demo = gr.Blocks()
def transcribe_speech(filepath):
    if filepath is None:
        gr.Warning("No text found, please retry.")
        return ""
    narrated_text=pipe(filepath)
    return narrated_text['sampling_rate'],narrated_text['audio']
    
mic_transcribe = gr.Interface(
    fn=transcribe_speech,
    inputs=gr.Textbox(label="Text",lines=3),
    outputs="audio",
    allow_flagging="never")

file_transcribe = gr.Interface(
    fn=transcribe_speech,
    inputs=gr.Audio(sources="upload",
                    type="filepath"),
    outputs="audio",
    #outputs=gr.Audio(label="Translated Message"),
    allow_flagging="never"
)
with demo:
    gr.TabbedInterface(
        [mic_transcribe],
        ["Transcribe Microphone"],
    )

demo.launch(share=True)
demo.close()