Update app.py

app.py

@@ -1,29 +1,11 @@
 import streamlit as st
 import torchaudio
 import torchaudio.transforms as T
-from transformers import pipeline, AutoProcessor, AutoModelForCTC, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForTextToSpectrogram
+from transformers import AutoProcessor, AutoModelForCTC, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForTextToSpectrogram
 import torch
 import numpy as np
 import io
 
-def transcribe_audio(audio_bytes):
-    # Load audio
-    waveform, sample_rate = torchaudio.load(io.BytesIO(audio_bytes), normalize=True)
-    
-    # Resample to 16kHz if necessary
-    if sample_rate != 16000:
-        resampler = T.Resample(orig_freq=sample_rate, new_freq=16000)
-        waveform = resampler(waveform)
-        sample_rate = 16000
-    
-    # Transcription
-    inputs = asr_processor(waveform.squeeze().numpy(), sampling_rate=sample_rate, return_tensors="pt", padding=True)
-    with torch.no_grad():
-        logits = asr_model(input_values=inputs.input_values).logits
-    predicted_ids = torch.argmax(logits, dim=-1)
-    transcription = asr_processor.decode(predicted_ids[0])
-    return transcription
-
 # Load models
 asr_model = AutoModelForCTC.from_pretrained("facebook/wav2vec2-large-robust-ft-swbd-300h")
 asr_processor = AutoProcessor.from_pretrained("facebook/wav2vec2-large-robust-ft-swbd-300h")
@@ -68,9 +50,20 @@ def generate_reply(text):
     return reply
 
 def text_to_speech(text):
+    # Load speaker embeddings
+    from transformers import Wav2Vec2ForCTC, Wav2Vec2Tokenizer
+    from datasets import load_dataset
+    
+    # Load pre-trained speaker embeddings (assuming you have downloaded them)
+    dataset = load_dataset("Matthijs/cmu-arctic-xvectors")
+    speaker_embeddings = dataset['train'][0]['xvector']
+    
     inputs = tts_processor(text=text, return_tensors="pt")
     with torch.no_grad():
-        spectrogram = tts_model.generate(**inputs)
+        spectrogram = tts_model.generate(
+            **inputs,
+            speaker_embeddings=speaker_embeddings
+        )
     return spectrogram
 
 # Streamlit app
@@ -98,9 +91,11 @@ if audio_input:
     
     # Convert text to speech
     spectrogram = text_to_speech(reply_text)
-    # Save spectrogram to file
+    
+    # Convert spectrogram to waveform for saving
+    waveform = tts_processor.convert_spectrogram_to_waveform(spectrogram)
     audio_file = io.BytesIO()
-    torchaudio.save(audio_file, spectrogram, 22050)  # assuming 22050 Hz sample rate
+    torchaudio.save(audio_file, waveform, 22050)  # assuming 22050 Hz sample rate
     audio_file.seek(0)
     
     st.audio(audio_file, format="audio/wav")