import gradio as gr
import torch
from tangoflux import TangoFluxInference

# Initialize model (auto-downloads on first run)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = TangoFluxInference(name="declare-lab/TangoFlux", device=device)

def generate_audio(prompt, steps, duration):
    audio = model.generate(prompt=prompt, steps=steps, duration=duration)
    return 44100, audio

# Gradio interface
demo = gr.Interface(
    fn=generate_audio,
    inputs=[
        gr.Textbox(lines=2, label="Prompt"),
        gr.Slider(25, 100, value=50, step=1, label="Steps"),
        gr.Slider(1, 30, value=10, step=1, label="Duration (s)")
    ],
    outputs=gr.Audio(type="numpy", label="Generated Audio"),
    title="TangoFlux: Text-to-Audio Generation",
    description="Generate sound effects (e.g., lightning, thunder) using TangoFlux."
)

demo.queue().launch()