import gradio as gr
import torch
from audioldm2.pipeline import build_model, text_to_audio

# Chargement du modèle AudioLDM2
device = "cuda" if torch.cuda.is_available() else "cpu"
ldm_model = build_model(device=device, model_name="audioldm2-full")

# Fonction appelée par Gradio
def generate_sound(prompt):
    if not prompt:
        return None, "⚠️ Entrez un mot ou une phrase"
    
    print(f"🔊 Génération à partir de : {prompt}")
    audios = text_to_audio(ldm_model, prompt, duration=3, guidance_scale=3.5, n_candidate_gen_per_text=1)
    audio_array = audios[0].cpu().numpy()
    return (16000, audio_array), f"✅ Généré depuis : {prompt}"

# Interface Gradio
with gr.Blocks() as demo:
    gr.Markdown("## 🎧 Générateur de sons avec AudioLDM2")
    
    with gr.Row():
        prompt_input = gr.Textbox(label="Texte ou mot-clé", placeholder="ex : explosion, footsteps, birds in jungle")
        generate_btn = gr.Button("Générer le son")
    
    output_audio = gr.Audio(label="Aperçu audio")
    status = gr.Textbox(label="Statut", interactive=False)

    generate_btn.click(fn=generate_sound, inputs=prompt_input, outputs=[output_audio, status])

demo.launch()