speech_conversation_app.py

import os
import time
import numpy as np
import torch
import gradio as gr
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, AutoProcessor, AutoModelForSpeechSeq2Seq
from datasets import load_dataset
import soundfile as sf

# Global variables to track latency
latency_ASR = 0.0
latency_LLM = 0.0
latency_TTS = 0.0

# Global variables to store conversation state
conversation_history = []
audio_output = None

# ASR Models
ASR_OPTIONS = {
    "Whisper Small": "openai/whisper-small",
    "Wav2Vec2": "facebook/wav2vec2-base-960h"
}

# LLM Models 
LLM_OPTIONS = {
    "Llama-2 7B Chat": "meta-llama/Llama-2-7b-chat-hf",
    "Flan-T5 Small": "google/flan-t5-small"
}

# TTS Models
TTS_OPTIONS = {
    "VITS": "espnet/kan-bayashi_ljspeech_vits", 
    "FastSpeech2": "espnet/kan-bayashi_ljspeech_fastspeech2"
}

# Load models
asr_models = {}
llm_models = {}
tts_models = {}

def load_asr_model(model_name):
    """Load ASR model from Hugging Face"""
    global asr_models
    
    if model_name not in asr_models:
        print(f"Loading ASR model: {model_name}")
        model_id = ASR_OPTIONS[model_name]
        
        if "whisper" in model_id:
            asr_models[model_name] = pipeline("automatic-speech-recognition", model=model_id)
        else:  # wav2vec2
            processor = AutoProcessor.from_pretrained(model_id)
            model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id)
            asr_models[model_name] = {"processor": processor, "model": model}
    
    return asr_models[model_name]

def load_llm_model(model_name):
    """Load LLM model from Hugging Face"""
    global llm_models
    
    if model_name not in llm_models:
        print(f"Loading LLM model: {model_name}")
        model_id = LLM_OPTIONS[model_name]
        
        tokenizer = AutoTokenizer.from_pretrained(model_id)
        model = AutoModelForCausalLM.from_pretrained(
            model_id, 
            torch_dtype=torch.float16, 
            device_map="auto"
        )
        
        llm_models[model_name] = {
            "model": model,
            "tokenizer": tokenizer
        }
    
    return llm_models[model_name]

def load_tts_model(model_name):
    """Load TTS model using ESPnet"""
    global tts_models
    
    if model_name not in tts_models:
        print(f"Loading TTS model: {model_name}")
        try:
            # Import ESPnet TTS modules
            from espnet2.bin.tts_inference import Text2Speech
            
            model_id = TTS_OPTIONS[model_name]
            tts = Text2Speech.from_pretrained(model_id)
            tts_models[model_name] = tts
            
        except ImportError:
            print("ESPnet not installed. Using mock TTS for demonstration.")
            tts_models[model_name] = "mock_tts"
    
    return tts_models[model_name]

def transcribe_audio(audio_data, sr, asr_model_name):
    """Transcribe audio using selected ASR model"""
    global latency_ASR
    
    start_time = time.time()
    
    model = load_asr_model(asr_model_name)
    
    if "whisper" in ASR_OPTIONS[asr_model_name]:
        result = model({"array": audio_data, "sampling_rate": sr})
        transcript = result["text"]
    else:  # wav2vec2
        inputs = model["processor"](audio_data, sampling_rate=sr, return_tensors="pt")
        with torch.no_grad():
            outputs = model["model"].generate(**inputs)
        transcript = model["processor"].batch_decode(outputs, skip_special_tokens=True)[0]
    
    latency_ASR = time.time() - start_time
    return transcript

def generate_response(transcript, llm_model_name, system_prompt):
    """Generate response using selected LLM model"""
    global latency_LLM, conversation_history
    
    start_time = time.time()
    
    model_info = load_llm_model(llm_model_name)
    model = model_info["model"]
    tokenizer = model_info["tokenizer"]
    
    # Format the prompt based on the model
    if "llama" in LLM_OPTIONS[llm_model_name].lower():
        # Format for Llama models
        if not conversation_history:
            conversation_history.append({"role": "system", "content": system_prompt})
        
        conversation_history.append({"role": "user", "content": transcript})
        
        prompt = tokenizer.apply_chat_template(
            conversation_history,
            tokenize=False,
            add_generation_prompt=True
        )
    else:
        # Format for T5 models
        prompt = f"{system_prompt}\nUser: {transcript}\nAssistant:"
    
    # Generate text
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(
            input_ids,
            max_new_tokens=100,
            temperature=0.7,
            top_p=0.9,
        )
    
    # Decode the response
    if "llama" in LLM_OPTIONS[llm_model_name].lower():
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        # Extract just the assistant's response
        response = response.split("Assistant: ")[-1].strip()
    else:
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Add to conversation history
    conversation_history.append({"role": "assistant", "content": response})
    
    latency_LLM = time.time() - start_time
    return response

def synthesize_speech(text, tts_model_name):
    """Synthesize speech using selected TTS model"""
    global latency_TTS
    
    start_time = time.time()
    
    tts = load_tts_model(tts_model_name)
    
    if tts == "mock_tts":
        # Mock TTS response for demonstration
        # In a real implementation, this would use the ESPnet model
        # Load a sample audio file for demonstration
        try:
            sample_rate = 16000
            # Generate a simple sine wave as demo audio
            duration = 2  # seconds
            t = np.linspace(0, duration, int(sample_rate * duration), endpoint=False)
            audio_data = 0.5 * np.sin(2 * np.pi * 220 * t)  # 220 Hz sine wave
        except Exception as e:
            print(f"Error generating mock audio: {e}")
            audio_data = np.zeros(16000)  # 1 second of silence
            sample_rate = 16000
    else:
        # Use actual ESPnet TTS model
        with torch.no_grad():
            wav = tts(text)["wav"]
        audio_data = wav.numpy()
        sample_rate = tts.fs
    
    latency_TTS = time.time() - start_time
    return (sample_rate, audio_data)

def process_speech(
    audio_input,
    asr_option,
    llm_option,
    tts_option,
    system_prompt
):
    """Process speech: ASR -> LLM -> TTS pipeline"""
    global audio_output
    
    # Check if audio input is available
    if audio_input is None:
        return None, "", "", None
    
    # Get audio data
    sr, audio_data = audio_input
    
    # ASR: Speech to text
    transcript = transcribe_audio(audio_data, sr, asr_option)
    
    # LLM: Generate response
    response = generate_response(transcript, llm_option, system_prompt)
    
    # TTS: Text to speech
    audio_output = synthesize_speech(response, tts_option)
    
    # Return results
    return audio_input, transcript, response, audio_output

def display_latency():
    """Display latency information"""
    return f"""
    ASR Latency: {latency_ASR:.2f} seconds
    LLM Latency: {latency_LLM:.2f} seconds
    TTS Latency: {latency_TTS:.2f} seconds
    Total Latency: {latency_ASR + latency_LLM + latency_TTS:.2f} seconds
    """

def reset_conversation():
    """Reset the conversation history"""
    global conversation_history, audio_output
    conversation_history = []
    audio_output = None
    return None, "", "", None, ""

# Create Gradio interface
with gr.Blocks(title="Conversational Speech System") as demo:
    gr.Markdown(
        """
        # Conversational Speech System with ASR, LLM, and TTS
        
        This demo showcases a complete speech-to-speech conversation system using:
        - **ASR** (Automatic Speech Recognition) to convert your speech to text
        - **LLM** (Large Language Model) to generate responses
        - **TTS** (Text-to-Speech) to convert the responses to speech
        
        Speak into your microphone and the system will respond with synthesized speech.
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            # Input components
            audio_input = gr.Audio(
                sources=["microphone"],
                type="numpy",
                label="Speak here",
            )
            
            system_prompt = gr.Textbox(
                label="System Prompt (instructions for the LLM)",
                value="You are a helpful and friendly AI assistant. Keep your responses concise and under 3 sentences."
            )
            
            asr_dropdown = gr.Dropdown(
                choices=list(ASR_OPTIONS.keys()),
                value=list(ASR_OPTIONS.keys())[0],
                label="Select ASR Model"
            )
            
            llm_dropdown = gr.Dropdown(
                choices=list(LLM_OPTIONS.keys()),
                value=list(LLM_OPTIONS.keys())[0],
                label="Select LLM Model"
            )
            
            tts_dropdown = gr.Dropdown(
                choices=list(TTS_OPTIONS.keys()),
                value=list(TTS_OPTIONS.keys())[0],
                label="Select TTS Model"
            )
            
            reset_btn = gr.Button("Reset Conversation")
            
        with gr.Column(scale=1):
            # Output components
            user_transcript = gr.Textbox(label="Your Speech (ASR Output)")
            system_response = gr.Textbox(label="AI Response (LLM Output)")
            audio_output_component = gr.Audio(label="AI Voice Response", autoplay=True)
            latency_info = gr.Textbox(label="Performance Metrics")
    
    # Set up event handlers
    audio_input.change(
        process_speech,
        inputs=[audio_input, asr_dropdown, llm_dropdown, tts_dropdown, system_prompt],
        outputs=[audio_input, user_transcript, system_response, audio_output_component]
    ).then(
        display_latency,
        inputs=[],
        outputs=[latency_info]
    )
    
    reset_btn.click(
        reset_conversation,
        inputs=[],
        outputs=[audio_input, user_transcript, system_response, audio_output_component, latency_info]
    )

# Launch the app
if __name__ == "__main__":
    demo.launch()