from diffusers import AudioLDMPipeline
import torch
import gradio as gr
from googletrans import Translator

if torch.cuda.is_available():
    device = "cuda"
    torch_dtype = torch.float16
else:
    device = "cpu"
    torch_dtype = torch.float32

repo_id = "cvssp/audioldm-m-full"
pipe = AudioLDMPipeline.from_pretrained(repo_id, torch_dtype=torch.float16)
pipe.unet = torch.compile(pipe.unet)

prompt = "Techno music with a strong, upbeat tempo and high melodic riffs"


def generate_sound(text): 
    print(text)
    text=translate_text(text)
    print(text)
    audio = pipe(text, num_inference_steps=10, audio_length_in_s=5.0).audios[0]
    rate =160000
    return rate, audio
def translate_text(text):
    translator = Translator()
    translated_text=translator.translate(text, src='es',dest="en")
    return translated_text.text

demo = gr.Blocks()
with demo:
    with gr.Row():
        with gr.Column():

            text = gr.Textbox(value="Ingrese el texto:")
            button = gr.Button(value="Generar")

        with gr.Column():
            output = gr.Audio()
    button.click(generate_sound,text,output)


demo.launch()