import gradio as gr
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor
import torch
import soundfile as sf
import tempfile

# Load model and processor
model_id = "jellecali8/somali_tts_model"
processor = AutoProcessor.from_pretrained(model_id)
model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id)

def tts(text):
    inputs = processor(text, return_tensors="pt")
    with torch.no_grad():
        outputs = model.generate(**inputs)
    audio = outputs[0].cpu().numpy()

    # Save to temporary file
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
        sf.write(f.name, audio, samplerate=16000)
        return f.name

iface = gr.Interface(
    fn=tts,
    inputs=gr.Textbox(lines=2, label="Enter Somali Text"),
    outputs=gr.Audio(label="Generated Speech"),
    title="Somali TTS Demo",
    description="Ku qor qoraalka Somali, kadib dhageyso codka."
)

iface.launch()