import gradio as gr

from inference import SAMPLE_RATE, generate_audio, load_model


pipeline, vocoder, device, dtype = load_model()


def predict(prompt, seed):
    audio = generate_audio(
        pipeline, vocoder, prompt, int(seed), device=device, dtype=dtype
    )
    return SAMPLE_RATE, audio


demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Rain and thunder"),
        gr.Number(label="Seed", value=42, precision=0),
    ],
    outputs=gr.Audio(label="Generated audio"),
    title="SwiftAudio",
    description="One-step text-to-audio generation with audio-free distillation.",
    examples=[
        ["Rain and thunder", 42],
        ["Ocean waves with seabirds", 42],
        ["A train whistles", 42],
    ],
)


if __name__ == "__main__":
    demo.launch()