import gradio as gr
import torch
from audiocraft.models import AudioGen
from audiocraft.data.audio import audio_write
import tempfile
import os

# Load the pre-trained model
model = AudioGen.get_pretrained('facebook/audiogen-medium')
model.set_generation_params(duration=5)

def generate_audio(text_prompt):
    try:
        wav = model.generate([text_prompt])[0]  # generates 1 sample.
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file:
            audio_write(temp_file.name, wav.cpu(), model.sample_rate, strategy="loudness", loudness_compressor=True)
            temp_filepath = temp_file.name
        return temp_filepath
    except Exception as e:
        return f"Error: {e}"

iface = gr.Interface(
    fn=generate_audio,
    inputs=gr.Textbox(lines=2, placeholder="Enter your text prompt here..."),
    outputs=gr.Audio(),
    title="AudioGen Gradio App",
    description="Generate audio from a text prompt using AudioCraft.",
)

iface.launch()