import spaces import os import shutil import time from typing import Generator, Optional, Tuple import gradio as gr import nltk import numpy as np import torch from huggingface_hub import HfApi import librosa from espnet2.sds.espnet_model import ESPnetSDSModelInterface # ------------------------ # Hyperparameters # ------------------------ access_token = os.environ.get("HF_TOKEN") ASR_name="espnet/owsm_ctc_v4_1B" LLM_name="Qwen/Qwen3-4B-Instruct-2507" TTS_name="espnet/kan-bayashi_ljspeech_vits" ASR_options="espnet/owsm_ctc_v4_1B".split(",") LLM_options="Qwen/Qwen3-4B-Instruct-2507".split(",") TTS_options="espnet/kan-bayashi_ljspeech_vits,espnet/kan-bayashi_jsut_full_band_vits_prosody,espnet/kan-bayashi_csmsc_full_band_vits".split(",") # Create display names with language labels for TTS models TTS_display_names = [ "English - espnet/kan-bayashi_ljspeech_vits", "Japanese - espnet/kan-bayashi_jsut_full_band_vits_prosody", "Chinese - espnet/kan-bayashi_csmsc_full_band_vits" ] TTS_value_map = {display: value for display, value in zip(TTS_display_names, TTS_options)} TTS_reverse_map = {value: display for display, value in zip(TTS_display_names, TTS_options)} Eval_options="Latency,TTS Intelligibility,TTS Speech Quality,ASR WER,Text Dialog Metrics" upload_to_hub=None dialogue_model = ESPnetSDSModelInterface( ASR_name, LLM_name, TTS_name, "Cascaded", access_token ) ASR_curr_name=None LLM_curr_name=None TTS_curr_name=None latency_ASR = 0.0 latency_LM = 0.0 latency_TTS = 0.0 LLM_response_arr = [] total_response_arr = [] enable_btn = gr.Button(interactive=True, visible=True) # ------------------------ # Function Definitions # ------------------------ def handle_eval_selection( option: str, TTS_audio_output: str, LLM_Output: str, ASR_audio_output: str, ASR_transcript: str, ): """ Handles the evaluation of a selected metric based on user input and provided outputs. This function evaluates different aspects of a casacaded conversational AI pipeline, such as: Latency, TTS intelligibility, TTS speech quality, ASR WER, and text dialog metrics. It is designed to integrate with Gradio via multiple yield statements, allowing updates to be displayed in real time. Parameters: ---------- option : str The evaluation metric selected by the user. Supported options include: - "Latency" - "TTS Intelligibility" - "TTS Speech Quality" - "ASR WER" - "Text Dialog Metrics" TTS_audio_output : np.ndarray The audio output generated by the TTS module for evaluation. LLM_Output : str The text output generated by the LLM module for evaluation. ASR_audio_output : np.ndarray The audio input/output used for ASR evaluation. ASR_transcript : str The transcript generated by the ASR module for evaluation. Returns: ------- str A string representation of the evaluation results. The specific result depends on the selected evaluation metric: - "Latency": Latencies of ASR, LLM, and TTS modules. - "TTS Intelligibility": A range of scores indicating how intelligible the TTS audio output is based on different reference ASR models. - "TTS Speech Quality": A range of scores representing the speech quality of the TTS audio output. - "ASR WER": The Word Error Rate (WER) of the ASR output based on different judge ASR models. - "Text Dialog Metrics": A combination of perplexity, diversity metrics, and relevance scores for the dialog. Raises: ------ ValueError If the `option` parameter does not match any supported evaluation metric. Example: ------- >>> result = handle_eval_selection( option="Latency", TTS_audio_output=audio_array, LLM_Output="Generated response", ASR_audio_output=audio_input, ASR_transcript="Expected transcript" ) >>> print(result) "ASR Latency: 0.14 LLM Latency: 0.42 TTS Latency: 0.21" """ global LLM_response_arr global total_response_arr return None def handle_eval_selection_E2E( option: str, TTS_audio_output: str, LLM_Output: str, ): """ Handles the evaluation of a selected metric based on user input and provided outputs. This function evaluates different aspects of an E2E conversational AI model, such as: Latency, TTS intelligibility, TTS speech quality, and text dialog metrics. It is designed to integrate with Gradio via multiple yield statements, allowing updates to be displayed in real time. Parameters: ---------- option : str The evaluation metric selected by the user. Supported options include: - "Latency" - "TTS Intelligibility" - "TTS Speech Quality" - "Text Dialog Metrics" TTS_audio_output : np.ndarray The audio output generated by the TTS module for evaluation. LLM_Output : str The text output generated by the LLM module for evaluation. Returns: ------- str A string representation of the evaluation results. The specific result depends on the selected evaluation metric: - "Latency": Latency of the entire system. - "TTS Intelligibility": A range of scores indicating how intelligible the TTS audio output is based on different reference ASR models. - "TTS Speech Quality": A range of scores representing the speech quality of the TTS audio output. - "Text Dialog Metrics": A combination of perplexity and diversity metrics for the dialog. Raises: ------ ValueError If the `option` parameter does not match any supported evaluation metric. Example: ------- >>> result = handle_eval_selection( option="Latency", TTS_audio_output=audio_array, LLM_Output="Generated response", ) >>> print(result) "Total Latency: 2.34" """ global LLM_response_arr global total_response_arr return def start_warmup(): """ Initializes and warms up the dialogue and evaluation model. This function is designed to ensure that all components of the dialogue model are pre-loaded and ready for execution, avoiding delays during runtime. """ global dialogue_model global ASR_options global LLM_options global TTS_options global ASR_name global LLM_name global TTS_name remove=0 for opt_count in range(len(ASR_options)): opt_count-=remove if opt_count>=len(ASR_options): break print(opt_count) print(ASR_options) opt = ASR_options[opt_count] try: for _ in dialogue_model.handle_ASR_selection(opt): continue except Exception as e: print(e) print("Removing " + opt + " from ASR options since it cannot be loaded.") ASR_options = ASR_options[:opt_count] + ASR_options[(opt_count + 1) :] remove+=1 if opt == ASR_name: ASR_name = ASR_options[0] for opt_count in range(len(LLM_options)): opt_count-=remove if opt_count>=len(LLM_options): break opt = LLM_options[opt_count] try: for _ in dialogue_model.handle_LLM_selection(opt): continue except Exception as e: print(e) print("Removing " + opt + " from LLM options since it cannot be loaded.") LLM_options = LLM_options[:opt_count] + LLM_options[(opt_count + 1) :] remove+=1 if opt == LLM_name: LLM_name = LLM_options[0] for opt_count in range(len(TTS_options)): opt_count-=remove if opt_count>=len(TTS_options): break opt = TTS_options[opt_count] try: for _ in dialogue_model.handle_TTS_selection(opt): continue except Exception as e: print(e) print("Removing " + opt + " from TTS options since it cannot be loaded.") TTS_options = TTS_options[:opt_count] + TTS_options[(opt_count + 1) :] remove+=1 if opt == TTS_name: TTS_name = TTS_options[0] # dialogue_model.handle_E2E_selection() dialogue_model.client = None for _ in dialogue_model.handle_TTS_selection(TTS_name): continue for _ in dialogue_model.handle_ASR_selection(ASR_name): continue for _ in dialogue_model.handle_LLM_selection(LLM_name): continue dummy_input = ( torch.randn( (3000), dtype=getattr(torch, "float16"), device="cpu", ) .cpu() .numpy() ) dummy_text = "This is dummy text" for opt in Eval_options: handle_eval_selection(opt, dummy_input, dummy_text, dummy_input, dummy_text) def flash_buttons(): """ Enables human feedback buttons after displaying system output. """ btn_updates = (enable_btn,) * 8 yield ( "", "", ) + btn_updates @spaces.GPU def handle_TTS_selection_wrapper(display_name: str): """ Wrapper function to convert TTS display name to actual model name before calling the dialogue model's handle_TTS_selection. This runs model-selection side effects only and intentionally does not update output widgets, so existing ASR/LLM/TTS outputs stay visible. Args: display_name: The display name of the TTS model (with language label) Returns: None """ # Convert display name to actual model name actual_name = TTS_value_map.get(display_name, display_name) # Run the original handler for side effects only. for _ in dialogue_model.handle_TTS_selection(actual_name): continue @spaces.GPU def handle_LLM_selection_wrapper(option: str): """ Run LLM model selection without clearing existing output widgets. """ for _ in dialogue_model.handle_LLM_selection(option): continue @spaces.GPU def handle_ASR_selection_wrapper(option: str): """ Run ASR model selection without clearing existing output widgets. """ for _ in dialogue_model.handle_ASR_selection(option): continue @spaces.GPU def process_audio_file( audio_file: Optional[Tuple[int, np.ndarray]], TTS_option: str, ASR_option: str, LLM_option: str, type_option: str, input_text: str, ): """ Processes a recorded audio file through the dialogue system. This function handles the transcription of an uploaded audio file and its transformation through a cascaded conversational AI system. It processes the entire audio file at once (offline mode). Args: audio_file: A tuple containing: - `sr`: Sample rate of the audio file. - `y`: Audio data array. TTS_option: Selected TTS model option (display name). ASR_option: Selected ASR model option. LLM_option: Selected LLM model option. type_option: Type of system ("Cascaded" or "E2E"). input_text: System prompt for the LLM. Returns: Tuple[str, str, Optional[Tuple[int, np.ndarray]]]: A tuple containing: - ASR output text (transcription). - Generated LLM output text (response). - Audio output as a tuple of sample rate and audio waveform (TTS). Notes: - Processes the complete audio file in one go. - Updates latency metrics. """ global latency_ASR global latency_LM global latency_TTS global LLM_response_arr global total_response_arr # Convert display name to actual model name TTS_option = TTS_value_map.get(TTS_option, TTS_option) if audio_file is None: gr.Info("Please upload an audio file.") return "", "", None # Ensure the selected models are actually loaded before processing for _ in dialogue_model.handle_TTS_selection(TTS_option): continue for _ in dialogue_model.handle_ASR_selection(ASR_option): continue for _ in dialogue_model.handle_LLM_selection(LLM_option): continue # Extract sample rate and audio data sr, y = audio_file # Convert stereo to mono if needed if len(y.shape) > 1: print(f"Converting stereo audio (shape: {y.shape}) to mono") y = np.mean(y, axis=1) # Ensure audio is float32 and normalized if y.dtype != np.float32: y = y.astype(np.float32) # Normalize if not already normalized if np.abs(y).max() > 1.0: y = y / np.abs(y).max() # Resample to 16000 Hz if needed (required for VAD and ASR) target_sr = 16000 if sr != target_sr: print(f"Resampling audio from {sr} Hz to {target_sr} Hz") y = librosa.resample(y, orig_sr=sr, target_sr=target_sr) sr = target_sr print(f"Processed audio: sr={sr}, shape={y.shape}, dtype={y.dtype}, range=[{y.min():.3f}, {y.max():.3f}]") # Initialize chat with system prompt dialogue_model.chat.init_chat( { "role": "system", "content": input_text, } ) # Initialize variables asr_output_str = "" text_str = "" audio_output = None audio_output1 = None stream = y # Use entire audio file as stream # Process the audio file ( asr_output_str, text_str, audio_output, audio_output1, latency_ASR, latency_LM, latency_TTS, stream, change, ) = dialogue_model( y, sr, stream, asr_output_str, text_str, audio_output, audio_output1, latency_ASR, latency_LM, latency_TTS, ) # Store results if change: print("Processing completed") if asr_output_str != "": total_response_arr.append(asr_output_str.replace("\n", " ")) LLM_response_arr.append(text_str.replace("\n", " ")) total_response_arr.append(text_str.replace("\n", " ")) return asr_output_str, text_str, audio_output # ------------------------ # Executable Script # ------------------------ api = HfApi() nltk.download("averaged_perceptron_tagger_eng") start_warmup() default_instruct=( "You are a helpful and friendly AI " "assistant. " "You are polite, respectful, and aim to " "provide concise and complete responses of " "less than 15 words." ) import pandas as pd # Usage examples: (label, prompt, TTS_display_name) # Use TTS_reverse_map[TTS_name] for default TTS; or a specific TTS_display_names entry example_translate_ja_prompt = "You are a translator. Translate what user says into Japanese." example_summarize_prompt = "Your task is to summarize what user says in one short sentence no more than 10 words." example_story_prompt = ( "You are an Imaginative Story Writer.\n\n" "Given user input, craft compelling fiction with:\n" "- strong scene setting\n" "- character voice\n" "- emotional arc\n" "- satisfying narrative progression\n\n" "Guidelines:\n" "- Respect user-provided anchors (characters, world, plot constraints).\n" "- Expand creatively between anchors with sensory detail and subtext.\n" "- Use \"show, don't tell\" where possible.\n" "- Maintain internal logic and continuity.\n" "- End scenes with momentum unless user requests a full ending.\n" "- Output only story content." ) example_audio_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "audio.wav") examples = [ ["Translation to Japanese", example_translate_ja_prompt, TTS_display_names[1]], ["Summarization (≤10 words)", example_summarize_prompt, TTS_reverse_map[TTS_name]], ["Imaginative Story Writer", example_story_prompt, TTS_reverse_map[TTS_name]], ] with gr.Blocks( title="ESPnet-SDS Offline Audio Processing", ) as demo: with gr.Row(): gr.Markdown( """ ## ESPnet Prompt Editing Demo Welcome to our prompt editing demo using ESPnet-SDS toolkit. **How to use:** 1. Upload or record an audio file 2. Configure the LLM prompt and select models 3. Click "Process Audio" to transcribe and generate a response The system will: - **Transcribe** your audio using ASR (Automatic Speech Recognition) - **Generate** a response using the selected LLM - **Synthesize** speech output using TTS (Text-to-Speech) **Click the examples in the bottom** for sample prompts and configurations. For more details, refer to the [README] (https://github.com/siddhu001/espnet/tree/sds_demo_recipe/egs2/TEMPLATE/sds1#how-to-use). """ ) type_radio = gr.State("Cascaded") with gr.Row(): with gr.Column(scale=1): user_audio = gr.Audio( sources=["upload", "microphone"], type="numpy", label="Upload or Record Audio File", format="wav", ) input_text = gr.Textbox( label="LLM prompt", visible=True, interactive=True, value=default_instruct, ) ASR_radio = gr.Radio( choices=ASR_options, label="Choose ASR:", value=ASR_name, ) LLM_radio = gr.Radio( choices=LLM_options, label="Choose LLM:", value=LLM_name, ) radio = gr.Radio( choices=TTS_display_names, label="Choose TTS:", value=TTS_reverse_map[TTS_name], ) E2Eradio = gr.Radio( choices=["mini-omni"], label="Choose E2E model:", value="mini-omni", visible=False, ) with gr.Column(scale=1): process_btn = gr.Button("Process Audio", variant="primary") output_asr_text = gr.Textbox(label="ASR Transcription", interactive=False) output_text = gr.Textbox(label="LLM Response", interactive=False) output_audio = gr.Audio(label="TTS Output", autoplay=True, visible=True, interactive=False) eval_radio = gr.Radio( choices=[ "Latency", "TTS Intelligibility", "TTS Speech Quality", "ASR WER", "Text Dialog Metrics", ], label="Choose Evaluation metrics:", visible=False, ) eval_radio_E2E = gr.Radio( choices=[ "Latency", "TTS Intelligibility", "TTS Speech Quality", "Text Dialog Metrics", ], label="Choose Evaluation metrics:", visible=False, ) output_eval_text = gr.Textbox(label="Evaluation Results", visible=False) gr.Examples( examples=[[row[1], row[2], example_audio_path] for row in examples], inputs=[input_text, radio, user_audio], label="Usage examples", example_labels=[row[0] for row in examples], ) natural_response = gr.Textbox( label="natural_response", visible=False, interactive=False ) diversity_response = gr.Textbox( label="diversity_response", visible=False, interactive=False ) ip_address = gr.Textbox(label="ip_address", visible=False, interactive=False) # Process button click event process_btn.click( process_audio_file, inputs=[user_audio, radio, ASR_radio, LLM_radio, type_radio, input_text], outputs=[output_asr_text, output_text, output_audio], ) radio.change( fn=handle_TTS_selection_wrapper, inputs=[radio], ) LLM_radio.change( fn=handle_LLM_selection_wrapper, inputs=[LLM_radio], ) ASR_radio.change( fn=handle_ASR_selection_wrapper, inputs=[ASR_radio], ) output_audio.play( flash_buttons, [], [natural_response, diversity_response] ) demo.queue(max_size=10, default_concurrency_limit=1) demo.launch(debug=True)