import torch
import gradio as gr
from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech
from datasets import load_dataset

# Load model and processor
processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts")

# Load a speaker embedding
embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
speaker_embedding = torch.tensor(embeddings_dataset[0]["xvector"]).unsqueeze(0)

# TTS Function
def text_to_speech(text):
    inputs = processor(text=text, return_tensors="pt")
    with torch.no_grad():
        speech = model.generate_speech(inputs["input_ids"], speaker_embedding)
    return speech.numpy()

# Set up Gradio interface
iface = gr.Interface(fn=text_to_speech, inputs="text", outputs="audio")
iface.launch()