app.py

import os
import gradio as gr
from pathlib import Path

from constants import STREAM_EVERY, APP_TMP_DIR
from ui import LED_DOT_OFF
from hf_dataset_utils import ALL_FILES

from stream_pipeline import (
    on_start_recording,
    on_stop_recording,
    shutdown_streamers,
    stream_step,
)
from offline_pipeline import (
    load_file_from_dataset,
    load_local_file,
    run_offline_pipeline,
)
from clean_up import purge_tmp_directory, cleanup_previous_run


_CSS_DIR = Path(__file__).resolve().parent / "assets"
with gr.Blocks() as demo:
    sample_stem = gr.State("")
    last_sample_stem = gr.State("")
    input_array = gr.State()
    streaming_sr = gr.State(None) 
    current_sample_rate = gr.State(None)
    
    gr.HTML(
        '<a href="https://ai-coustics.com/" target="_blank">'
        '<img src="https://mintcdn.com/ai-coustics/Sxcrv8jVSE2qWMR1/logo/dark.svg?fit=max&auto=format&n=Sxcrv8jVSE2qWMR1&q=85&s=7f26caaf21e963912961cbd8541e6d84" '
        'alt="ai-coustics Logo" width="400" style="display: block; margin: 0 auto;">'
        "</a>"
    )
    gr.Markdown(open("docs/intro.md", "r", encoding="utf-8").read())

    # ✅ Global controls (shared by both tabs)
    stt_model = gr.Radio(
        label="STT Model",
        choices=["Soniox STT-RT v3", "Deepgram Nova-3 RT"],
        value="Soniox STT-RT v3",
        interactive=True,
    )
    # Streamers are created on first start_recording to avoid Soniox "No audio received" timeout
    enhancement_level = gr.Slider(
        minimum=0,
        maximum=100,
        step=1,
        value=80,
        label="Enhancement level (%)",
        scale=2,
    )

    with gr.Tabs():
        with gr.Tab("Stream audio in real time") as stream_tab:
            with gr.Row(equal_height=False, elem_classes="stream-layout"):
                with gr.Column(scale=4, min_width=320):
                    with gr.Group(elem_classes="panel section-panel"):
                        gr.Markdown("### Input", elem_classes="title")
                        gr.Markdown(open("docs/online.md", "r", encoding="utf-8").read(), elem_classes="tab-description")

                        input_gain_db = gr.Slider(
                            minimum=0,
                            maximum=20,
                            step=0.5,
                            value=0,
                            label="Input gain (dB)",
                        )

                        audio_stream = gr.Audio(
                            sources=["microphone"],
                            streaming=True,
                            elem_id="audio_stream",
                        )

                with gr.Column(scale=6, min_width=420):
                    with gr.Group(elem_classes="status-panel output-panel"):
                            gr.Markdown("### Live Output", elem_classes="title")
                            with gr.Row(equal_height=True, elem_classes="status-indicators"):
                                with gr.Group(elem_classes="status-card"):
                                    gr.Markdown("**System Status**", elem_classes="status-card__label")
                                    system_status_led = gr.HTML(value=LED_DOT_OFF, show_label=False)
                                    system_status_text = gr.Markdown(
                                        value="Off",
                                        elem_classes="status-card__subtext",
                                    )

                                with gr.Group(elem_classes="status-card"):
                                    gr.Markdown("**Voice Activity**", elem_classes="status-card__label")
                                    vad_led = gr.HTML(value=LED_DOT_OFF, show_label=False)

                            with gr.Row(equal_height=True, elem_classes="transcript-row transcript-row--large"):
                                enhanced_text = gr.Textbox(
                                    label="Enhanced Transcript",
                                    lines=10,
                                    autoscroll=False,
                                )
                                raw_text = gr.Textbox(
                                    label="Raw Transcript",
                                    lines=10,
                                    autoscroll=False,
                                )

                        
        with gr.Tab("Dataset: Dawn Chorus") as dataset_tab:
            with gr.Group(elem_classes="panel section-panel"):
                gr.Markdown("### Input", elem_classes="title")
                gr.Markdown(open("docs/dawn_chorus.md", "r", encoding="utf-8").read(), elem_classes="tab-description")

                dataset_dropdown = gr.Dropdown(
                    choices=ALL_FILES, value="en_00352_i_h_33", label="Sample"
                )
                audio_file_from_dataset = gr.Audio(
                    type="filepath", interactive=False, buttons=["download"], autoplay=False
                )

        with gr.Tab("Upload local file") as upload_tab:
            with gr.Group(elem_classes="panel section-panel"):
                gr.Markdown("### Input", elem_classes="title")
                gr.Markdown(open("docs/local_file.md", "r", encoding="utf-8").read(), elem_classes="tab-description")

                audio_file_upload = gr.File(
                    file_types=[".wav", ".mp3", ".flac", ".m4a", ".ogg"],
                    file_count="single",
                    scale=3,
                )
                normalize = gr.Checkbox(label="Normalize audio", value=True)
                audio_preview = gr.Audio(
                    label="Preview",
                    autoplay=False,
                    interactive=False,
                )

    enhance_btn = gr.Button("Enhance with Quail Voice Focus 2.0", scale=2, visible=False)

    with gr.Group(elem_classes="panel results-card", visible=False) as results_card:
        result_title = gr.Markdown("", elem_classes="title")
        enhanced_audio = gr.Audio(
            type="numpy",
            interactive=False,
            buttons=["download"],
            autoplay=False,
        )

        with gr.Row(equal_height=True, elem_classes="results-row"):
            with gr.Column(scale=5, min_width=320, elem_classes="results-left"):
                noisy_image = gr.Image(
                    label="Input spectrogram", format="png", type="filepath"
                )
                enhanced_image = gr.Image(
                    label="Enhanced spectrogram", format="png", type="filepath"
                )

            with gr.Column(scale=5, min_width=320, elem_classes="results-right"):
                original_transcript = gr.Textbox(
                    label="Groundtruth transcription", lines=3, interactive=False
                )
                noisy_transcript = gr.Textbox(
                    label="Noisy transcription", lines=3, interactive=False
                )
                enhanced_transcript = gr.Textbox(
                    label="Transcription post enhancement", lines=3, interactive=False
                )

    # ------------------------------------------------------
    # TAB CHANGES
    # ------------------------------------------------------
    stream_tab.select(
        lambda: (
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(sources=["microphone"], streaming=True, interactive= True),
        ),
        outputs=[results_card, enhance_btn, audio_stream],
    )

    def _on_not_streaming_tab():
        shutdown_streamers()
        return (
            gr.update(streaming=False, interactive=False),
            gr.update(visible=True),
            LED_DOT_OFF,
            LED_DOT_OFF,
            "Off",
        )

    upload_tab.select(
        _on_not_streaming_tab,
        outputs=[audio_stream, enhance_btn, vad_led, system_status_led, system_status_text],  
    ).then(
        load_local_file,
        inputs=[audio_file_upload, normalize],
        outputs=[input_array, sample_stem, audio_preview, current_sample_rate],
    )
    
    dataset_tab.select(
        _on_not_streaming_tab,
        outputs=[audio_stream, enhance_btn, vad_led, system_status_led, system_status_text],
    ).then(
        load_file_from_dataset,
        inputs=dataset_dropdown,
        outputs=[audio_file_from_dataset, input_array, sample_stem, current_sample_rate],
    )
    # ------------------------------------------------------
    # STREAMING EVENTS
    # ------------------------------------------------------
    audio_stream.stream(
        stream_step,
        inputs=[audio_stream, streaming_sr, stt_model, enhancement_level, input_gain_db],
        outputs=[streaming_sr, system_status_led, system_status_text,enhanced_text, raw_text, vad_led],
        stream_every=STREAM_EVERY,
        time_limit=60 * 2,
        concurrency_limit=1,    
    )

    audio_stream.stop_recording(
        on_stop_recording,
        outputs=[vad_led, system_status_led, system_status_text, streaming_sr],
    ).then(
        shutdown_streamers,
    )

    audio_stream.start_recording(
        on_start_recording,  
        outputs=[enhanced_text, raw_text, system_status_led, system_status_text],
    )    

    # ------------------------------------------------------
    # OFFLINE EVENTS (DATASET + LOCAL FILE)
    # ------------------------------------------------------

    # Dataset dropdown selection triggers loading the audio file and hiding results until enhancement
    dataset_dropdown.change(
        lambda: gr.update(visible=False),
        outputs=results_card,
    ).then(
        load_file_from_dataset,
        inputs=dataset_dropdown,
        outputs=[audio_file_from_dataset, input_array, sample_stem, current_sample_rate],
    )

    # Uploading a local file triggers loading the audio file and hiding results until enhancement
    audio_file_upload.change(
        lambda: gr.update(visible=False),
        outputs=results_card,
    ).then(
        load_local_file,
        inputs=[audio_file_upload, normalize],
        outputs=[input_array, sample_stem, audio_preview, current_sample_rate],
    )
 
    normalize.change(
        load_local_file,
        inputs=[audio_file_upload, normalize],
        outputs=[input_array, sample_stem, audio_preview, current_sample_rate],
    )

    # Enhancement button: run pipeline with live transcript progress (dataset + local file modes).
    enhance_btn.click(
        cleanup_previous_run,
        inputs=[last_sample_stem]
    ).then(
        lambda: gr.update(visible=True),
        outputs=results_card,
    ).then(
        run_offline_pipeline,
        inputs=[input_array, current_sample_rate, enhancement_level, stt_model, sample_stem],
        outputs=[
            enhanced_audio,
            enhanced_image,
            noisy_image,
            original_transcript,
            noisy_transcript,
            enhanced_transcript,
            last_sample_stem,
            ]
    ).failure(
        lambda: gr.Warning("Enhancement failed. Please refresh page and make sure you have a stable connection.")
    )

os.makedirs(APP_TMP_DIR, exist_ok=True)
purge_tmp_directory(max_age_minutes=0, tmp_dir=APP_TMP_DIR)
demo.queue()
demo.launch(
    css=(_CSS_DIR / "styling.css").read_text(encoding="utf-8"),
    allowed_paths=[APP_TMP_DIR, "/tmp", "/"],
)