import streamlit as st
from transformers import pipeline, SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan
from datasets import load_dataset
import torch
import soundfile as sf
import os

# Function to generate speech using the pipeline method
def generate_speech_pipeline(text, speaker_embedding):
    synthesiser = pipeline("text-to-speech", "microsoft/speecht5_tts")
    speech = synthesiser(text, forward_params={"speaker_embeddings": speaker_embedding})
    return speech["audio"], speech["sampling_rate"]

# Function to generate speech using the processor + generate method
def generate_speech_processor(text, speaker_embedding):
    processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
    model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts")
    vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")

    inputs = processor(text=text, return_tensors="pt")
    speech = model.generate_speech(inputs["input_ids"], speaker_embedding, vocoder=vocoder)
    return speech.numpy(), 16000

def main():
    st.title("Text-to-Speech with SpeechT5")
    
    st.write("Enter the text you want to convert to speech:")

    text = st.text_area("Text", "Hello, my dog is cooler than you!")
    
    if st.button("Generate Speech"):
        st.write("Generating speech...")

        # Load xvector containing speaker's voice characteristics from a dataset
        embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
        speaker_embedding = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0)

        # Choose the method to generate speech
        method = st.selectbox("Choose the method for generating speech", ["Pipeline", "Processor + Generate"])

        if method == "Pipeline":
            audio, samplerate = generate_speech_pipeline(text, speaker_embedding)
        else:
            audio, samplerate = generate_speech_processor(text, speaker_embedding)

        # Save and play the generated speech
        output_path = "speech.wav"
        sf.write(output_path, audio, samplerate=samplerate)
        st.audio(output_path)

if __name__ == "__main__":
    main()