Update app.py

app.py

@@ -1,29 +1,54 @@
 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 spaces
 import os
-import numpy as np
-import re
-from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan
-from speechbrain.pretrained import EncoderClassifier
-from datasets import load_dataset
 
-# Load model directly
-from transformers import AutoProcessor, AutoModelForTextToSpectrogram
+# Load models and processors
+processor = AutoProcessor.from_pretrained("speecht5_finetuned_Aumkesh_tr")
+model = AutoModelForTextToSpectrogram.from_pretrained("speecht5_finetuned_Aumkesh_tr")
+vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")
 
-processor = AutoProcessor.from_pretrained("Aumkeshchy2003/speecht5_finetuned_Aumkesh_tr")
-model = AutoModelForTextToSpectrogram.from_pretrained("Aumkeshchy2003/speecht5_finetuned_Aumkesh_tr")
+# 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 English text to convert to speech")
-    ],
-    outputs=[
-        gr.Audio(label="Generated Speech", type="numpy")
-    ],
-    title="English SpeechT5 Text-to-Speech Demo",
-    description="Enter English text, and listen to the generated speech."
+    inputs=gr.Textbox(label="Enter the text"),
+    outputs=gr.Audio(label="Generated Speech"),
+    title="Text-to-Speech Converter",
+    description="Convert text to speech using the SpeechT5 model."
 )
-iface.launch(share=True)
+
+# Launch the app
+iface.launch()