import torch
import gradio as gr
from dia.model import Dia

# Use CPU
device = torch.device("cpu")

# Load the model once
model = Dia.from_pretrained(
    "nari-labs/Dia-1.6B-0626", compute_dtype="float32", device=device
)

def generate_audio(text):
    output = model.generate(text, use_torch_compile=False, verbose=False)
    output_path = "output.mp3"
    model.save_audio(output_path, output)
    return output_path

# Gradio UI
iface = gr.Interface(
    fn=generate_audio,
    inputs=gr.Textbox(lines=4, placeholder="Enter dialogue here..."),
    outputs=gr.Audio(type="filepath"),
    title="Dia - Text to Dialogue",
    description="Enter dialogue formatted like [S1] Hello. [S2] Hi there!",
)

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