fix2

app.py

@@ -1,7 +1,8 @@
 import gradio as gr
 import numpy as np
 import torch
-from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5ForSpeechToText
+from transformers import pipeline, SpeechT5Processor, SpeechT5ForTextToSpeech
+from datasets import load_dataset
 import soundfile as sf
 import tempfile
 import os
@@ -9,36 +10,41 @@ import os
 # Check if CUDA is available, otherwise use CPU
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load SpeechT5 models and processor
-processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_asr")
-asr_model = SpeechT5ForSpeechToText.from_pretrained("microsoft/speecht5_asr").to(device)
+# Load Whisper for ASR (much more reliable than SpeechT5 for ASR)
+print("Loading ASR model...")
+asr_pipeline = pipeline("automatic-speech-recognition", model="openai/whisper-small", device=device)
+
+# Load SpeechT5 for TTS
+print("Loading TTS model...")
+tts_processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
 tts_model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts").to(device)
 
-# Function to convert speech to text
-def speech_to_text(audio_dict):
-    # Extract the audio array from the dictionary
-    audio_array = audio_dict["array"]
-    
-    # Pass the audio array directly to the processor
-    inputs = processor(audio=audio_array, sampling_rate=16000, return_tensors="pt").input_values.to(device)
-    
-    with torch.no_grad():
-        logits = asr_model(inputs).logits
+# Load speaker embeddings for TTS
+print("Loading speaker embeddings...")
+embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
+speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0).to(device)
+
+# Function to convert speech to text using Whisper
+def speech_to_text(audio_data, sample_rate):
+    # Normalize audio data
+    audio_data = audio_data.flatten().astype(np.float32) / 32768.0
     
-    predicted_ids = torch.argmax(logits, dim=-1)
-    transcription = processor.batch_decode(predicted_ids)[0]
-    return transcription
+    # Process with Whisper
+    result = asr_pipeline({"raw": audio_data, "sampling_rate": sample_rate})
+    return result["text"]
 
-# Function to convert text to speech
+# Function to convert text to speech using SpeechT5
 def text_to_speech(text):
-    inputs = processor(text=text, return_tensors="pt").input_ids.to(device)
-    # Create dummy decoder input IDs (this is a simplification)
-    decoder_input_ids = torch.zeros((1, 1), dtype=torch.long).to(device)
+    # Process text input
+    inputs = tts_processor(text=text, return_tensors="pt").to(device)
+    
+    # Generate speech with speaker embeddings
     with torch.no_grad():
         speech = tts_model.generate_speech(
-            inputs,
-            decoder_input_ids=decoder_input_ids
+            inputs["input_ids"], 
+            speaker_embeddings=speaker_embeddings
         )
+    
     return speech
 
 # Gradio demo
@@ -55,12 +61,11 @@ def demo():
             if audio is None:
                 return None, "No audio detected."
             
-            # Convert audio to the correct format
+            # Get audio data
             sample_rate, audio_data = audio
-            audio_data = audio_data.flatten().astype(np.float32) / 32768.0  # Normalize to [-1.0, 1.0]
             
             # Speech-to-text
-            transcript = speech_to_text({"array": audio_data, "sampling_rate": sample_rate})
+            transcript = speech_to_text(audio_data, sample_rate)
             print(f"Transcribed: {transcript}")
             
             # Generate response (for simplicity, echo the transcript)