Update app.py

app.py

@@ -1,24 +1,51 @@
-import torch
 import gradio as gr
-from transformers import pipeline
+import torch
+from transformers import WhisperForConditionalGeneration, WhisperProcessor, AutoTokenizer, AutoModelForCausalLM, VitsModel
+
+# Load the models and tokenizers
+asr_model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny")
+asr_tokenizer = AutoTokenizer.from_pretrained("openai/whisper-tiny")
+asr_processor = WhisperProcessor.from_pretrained("openai/whisper-tiny")
+
+llm_model = AutoModelForCausalLM.from_pretrained("microsoft/Phi-3-mini-4k-instruct", trust_remote_code=True)
+llm_tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-4k-instruct")
+
+tts_model = VitsModel.from_pretrained("kakao-enterprise/vits-ljs")
+tts_tokenizer = AutoTokenizer.from_pretrained("kakao-enterprise/vits-ljs")
 
-asr_pipe = pipeline("automatic-speech-recognition", model="openai/whisper-tiny")
-llm_pipe = pipeline("text-generation", model="microsoft/Phi-3-mini-4k-instruct", trust_remote_code=True)
-tts_pipe = pipeline("text-to-speech", model="kakao-enterprise/vits-ljs")
+# Define the function to transcribe audio
+def transcribe_audio(audio_input):
+    processor = WhisperProcessor.from_pretrained("openai/whisper-tiny")
+    model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny")
+    model.config.forced_decoder_ids = None
+    
+    input_features = processor(audio_input, return_tensors="pt").input_features
+    predicted_ids = model.generate(input_features)
+    transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)
+    
+    return transcription
 
-def process_audio(audio):
-    # Convert audio to text using ASR pipeline
-    transcribed_text = asr_pipe(audio)[0]['transcription']
+# Define the function to generate text
+def generate_text(transcription):
+    input_ids = llm_tokenizer.encode(transcription, return_tensors="pt")
+    output_ids = llm_model.generate(input_ids)
+    generated_text = llm_tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    return generated_text
 
-    # Generate text using LLM pipeline
-    generated_text = llm_pipe(transcribed_text)[0]['generated_text']
+def text_to_speech(generated_text):
+    text_input = tts_tokenizer(generated_text, return_tensors="pt")
+    
+    with torch.no_grad():
+        output = tts_model(**text_input).waveform
+        
+    return output
 
-    # Convert generated text to speech using TTS pipeline
-    speech = tts_pipe(generated_text)[0]['audio']
+def my_alexa(audio_input):
+    transcription = transcribe_audio(audio_input)
+    generated_text = generate_text(transcription)
+    speech_output = text_to_speech(generated_text)
+    return speech_output
 
-    return speech
 
-audio_input = gr.Audio(label="Audio Input", sources="microphone")
-my_alexa_output = gr.Interface(fn=process_audio, inputs=audio_input, outputs="audio", title="MyAlexa AI Assistant", 
-                      description="An AI Assistant That Will Answer Your Questions!")
-my_alexa_output.launch(inbrowser=True)
+gr.Interface(fn=transcribe_audio, inputs=gr.Audio(sources="microphone"), outputs="text", title="MyAlexa AI Assistant",
+             description="An AI Assistant That Can Answer Your Questions!").launch(inbrowser=True, debug=True)