import gradio as gr
from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan
from transformers import AutoProcessor, AutoModelForTextToSpectrogram
from datasets import load_dataset
import torch
import soundfile as sf
import os

# Load model directly
processor = AutoProcessor.from_pretrained("Aumkeshchy2003/speecht5_finetuned_Aumkesh_English_tts")
model = AutoModelForTextToSpectrogram.from_pretrained("Aumkeshchy2003/speecht5_finetuned_Aumkesh_English_tts")
vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")

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

# Quantize the models
def quantize_model(model):
    quantized_model = torch.quantization.quantize_dynamic(
        model, {torch.nn.Linear}, dtype=torch.qint8
    )
    return quantized_model

# Only quantize the vocoder, as the main model might not be compatible
vocoder = quantize_model(vocoder)

# Move models to GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
vocoder = vocoder.to(device)
speaker_embeddings = speaker_embeddings.to(device)

# Use inference mode for faster computation
@torch.inference_mode()
def text_to_speech(text):
    inputs = processor(text=text, return_tensors="pt").to(device)
    speech = model.generate_speech(inputs["input_ids"], speaker_embeddings, vocoder=vocoder)
    speech = speech.cpu()  # Move back to CPU for saving
    output_path = "output.wav"
    sf.write(output_path, speech.numpy(), samplerate=16000)
    return output_path

# Create Gradio interface
iface = gr.Interface(
    fn=text_to_speech,
    inputs=gr.Textbox(label="Enter the text"),
    outputs=gr.Audio(label="Generated Speech"),
    title="Text-to-Speech",
    description="Convert text to speech."
)

# Launch the app
iface.launch()