app.py

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)