app.py

import os
import threading
import time
import gradio as gr
from pathlib import Path

from constants import STREAM_EVERY, WARMUP_SECONDS, APP_TMP_DIR
from ui import ACTIVE_LIGHT_GRAY, ACTIVE_LIGHT_RED, VAD_OFF_HTML
from hf_dataset_utils import ALL_FILES, get_transcript

from stream_pipeline import (
    clear_ui,
    get_live_transcripts,
    on_stop_recording,
    set_stt_streamer,
    stop_online_backend,
    transcribe_stream,
    shutdown_streamers,
)
from offline_pipeline import (
    load_file_from_dataset,
    load_local_file,
    run_offline_pipeline_streaming,
)
from utils import spec_image
from clean_up import purge_tmp_directory, cleanup_previous_run

WARMUP_TICKS = max(1, int(WARMUP_SECONDS * 2))  # timer ticks every 0.5s


def warmup_tick(on_stream_tab, warmup_elapsed, _current_html):
    if not on_stream_tab:
        return 0, ACTIVE_LIGHT_GRAY, gr.update(active=True)
    if warmup_elapsed >= WARMUP_TICKS:
        return warmup_elapsed, ACTIVE_LIGHT_RED, gr.update(active=False)
    return warmup_elapsed + 1, ACTIVE_LIGHT_GRAY, gr.update(active=True)


def process_with_live_transcript(
    input_array,
    enhancement_level,
    sample_stem,
    stt_model,
    last_sample_stem,
    current_sample_rate,
):  
    """Generator that runs the offline pipeline in real time (chunked): enhanced audio and
    both transcripts stream from the first chunk so playback and transcription start immediately."""
    progress_state = {}
    result_holder = {}

    def worker():
        try:
            result_holder["result"] = run_offline_pipeline_streaming(
                input_array,
                current_sample_rate,
                enhancement_level,
                sample_stem,
                stt_model,
                progress_state
            )
        except Exception as e:
            result_holder["error"] = e

    # 1) First yield: ground truth + input spectrogram only (no audio, no enhanced spec, no transcripts yet)
    _ = cleanup_previous_run(last_sample_stem)
    noisy_spec_path = f"{APP_TMP_DIR}/{sample_stem}_noisy_spectrogram.png"
    if input_array is not None:
        try:
            spec_image(input_array, sr = current_sample_rate).save(noisy_spec_path)
        except Exception:
            noisy_spec_path = None
    else:
        noisy_spec_path = None
    try:
        original_transcript = get_transcript(sample_stem)
    except Exception:
        original_transcript = "Unavailable"

    yield (
        gr.update(visible=True),
        None,                # enhanced_audio: set only in final yield (smooth playback)
        gr.update(value=None),  # enhanced_image: clear until step 3 (last)
        gr.update(value=noisy_spec_path),  # noisy_image: input spectrogram (step 1)
        original_transcript,
        "",
        "",
        sample_stem,
        None,
        "",
    )
    # Let the UI render step 1 before we flood with polling updates
    time.sleep(0.2)

    thread = threading.Thread(target=worker, daemon=True)
    thread.start()

    poll_interval = 0.05
    while "result" not in result_holder and "error" not in result_holder:
        time.sleep(poll_interval)
        # 2) Realtime: stream transcripts only; audio set in final yield for smooth playback
        yield (
            gr.update(visible=True),
            gr.update(),  # enhanced_audio: set only in final yield, then autoplay
            gr.update(),  # enhanced_image: reveal only in step 3 (final yield)
            gr.update(),  # noisy_image already set in step 1
            gr.update(),  # original_transcript unchanged
            gr.update(value=progress_state.get("noisy", "")),
            gr.update(value=progress_state.get("enhanced", "")),
            gr.update(),
            gr.update(),
            gr.update(),
        )

    if "error" in result_holder:
        raise result_holder["error"]

    (
        enhanced_spec_path,
        enhanced_transcript,
        noisy_transcript_with_wer,
        enhanced_audio,
        vad_labels,
        last_stem,
        enhanced_array,
        precomputed_noisy,
    ) = result_holder["result"]
    # 3) Last: reveal enhanced spectrogram (and final audio/transcripts)
    yield (
        gr.update(visible=True),
        gr.update(value=enhanced_audio, subtitles=vad_labels),
        enhanced_spec_path,   # enhanced_image: show only now
        noisy_spec_path,
        original_transcript,
        noisy_transcript_with_wer,
        enhanced_transcript,
        last_stem,
        enhanced_array,
        precomputed_noisy,
    )


_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()
    enhanced_array = gr.State()
    precomputed_noisy_transcript = gr.State("")
    current_sample_rate = gr.State(16000)  # default sample rate for dataset samples; updated on local file load if different

    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.Group(elem_classes="panel"):
        with gr.Tab("Stream audio in real time") as stream_tab:
            gr.Markdown(open("docs/online.md", "r", encoding="utf-8").read())
            with gr.Group(elem_classes="panel"):
                streaming_sr = gr.State(16000)  
                on_stream_tab = gr.State(True)  # True on load: stream tab is the default first tab
                warmup_elapsed = gr.State(0)
                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.Group(elem_classes="panel"):
                    with gr.Row(scale=1, elem_classes="status-row"):
                        active_light = gr.HTML(value=ACTIVE_LIGHT_GRAY, label="System Status", show_label=True)
                        vad_led = gr.HTML(value=VAD_OFF_HTML, label="Voice Activity Detection", show_label=True)
                    with gr.Column(scale=5, min_width=320):
                        enhanced_text = gr.Textbox(
                            label="Enhanced Transcribed Text", lines=6, autoscroll=False
                        )
                    with gr.Column(scale=5, min_width=320):
                        raw_text = gr.Textbox(
                            label="Raw Transcribed Text", lines=6, autoscroll=False
                        )
                # Poll transcript globals so interim results show live (streamers update in background)
                transcript_timer = gr.Timer(0.1, active=True)
                transcript_timer.tick(
                    get_live_transcripts,
                    inputs=None,
                    outputs=[enhanced_text, raw_text],
                    show_progress="hidden",
                )
                warmup_timer = gr.Timer(0.5, active=True)
                warmup_timer.tick(
                    warmup_tick,
                    inputs=[on_stream_tab, warmup_elapsed, active_light],
                    outputs=[warmup_elapsed, active_light, warmup_timer],
                    show_progress="hidden",
                )

        with gr.Tab("Dataset: Dawn Chorus") as dataset_tab:
            with gr.Row():
                gr.Markdown(open("docs/dawn_chorus.md", "r", encoding="utf-8").read())
            dataset_dropdown = gr.Dropdown(
                choices=ALL_FILES, value="en_00412_i_h_36", 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.Row():
                gr.Markdown(open("docs/local_file.md", "r", encoding="utf-8").read())
            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("")
        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
                )

    # ------------------------------------------------------
    # GLOBAL EVENTS (TAB CHANGES, MODEL CHANGES)
    # ------------------------------------------------------
    DEFAULT_DATASET_SAMPLE = "en_00412_i_h_36"

    def load_dataset_sample_on_tab_visit(dropdown_value):
        """Load the selected sample when visiting the Dataset tab; use default if dropdown is empty."""
        sample_id = dropdown_value or DEFAULT_DATASET_SAMPLE
        audio_path, arr, stem, sample_rate = load_file_from_dataset(sample_id)
        return sample_id, audio_path, arr, stem, sample_rate

    stream_tab.select(
        lambda: (
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(sources=["microphone"], streaming=True, interactive= True),
            True,
            0,
            ACTIVE_LIGHT_GRAY,
            gr.update(active=True),
        ),
        inputs=None,
        outputs=[results_card, enhance_btn, audio_stream, on_stream_tab, warmup_elapsed, active_light, warmup_timer],
    )

    upload_tab.select(
        lambda: (
            gr.update(streaming=False, interactive=False),
            gr.update(visible=True),
            *on_stop_recording(),
        ),
        inputs=None,
        outputs=[audio_stream,enhance_btn, vad_led, streaming_sr],
    )
    
    dataset_tab.select(
        lambda: ( 
            gr.update(streaming=False, interactive=False),
            gr.update(visible=True),
            *on_stop_recording(),
        ),
        outputs=[audio_stream,enhance_btn, vad_led, streaming_sr],
    ).then(
        load_dataset_sample_on_tab_visit,
        inputs=[dataset_dropdown],
        outputs=[dataset_dropdown, audio_file_from_dataset, input_array, sample_stem, current_sample_rate],
    )

    stt_model.change(
        fn=shutdown_streamers,
    ).then(
        clear_ui,
        inputs=None,
        outputs=[enhanced_text, raw_text],
    ).then(
        set_stt_streamer,
        inputs=[stt_model, streaming_sr],
        outputs=None,
    )

    # ------------------------------------------------------
    # ONLINE PIPELINE EVENTS (DATASET + UPLOAD TABS)
    # ------------------------------------------------------

    audio_stream.stream(
        fn=transcribe_stream,
        inputs=[streaming_sr, audio_stream, enhancement_level, input_gain_db, stt_model],
        outputs=[streaming_sr, 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, streaming_sr]
    )

    audio_stream.start_recording(
        clear_ui,
        inputs=None,
        outputs=[enhanced_text, raw_text],
    )
    

    # ------------------------------------------------------
    # OFFLINE PIPELINE EVENTS (DATASET + UPLOAD TABS)
    # ------------------------------------------------------

    # Dataset dropdown selection triggers loading the audio file and hiding results until enhancement
    dataset_dropdown.change(
        lambda: gr.update(visible=False),
        inputs=None,
        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(
        process_with_live_transcript,
        inputs=[input_array, enhancement_level, sample_stem, stt_model, last_sample_stem, current_sample_rate],
        outputs=[
            results_card,
            enhanced_audio,
            enhanced_image,
            noisy_image,
            original_transcript,
            noisy_transcript,
            enhanced_transcript,
            last_sample_stem,
            enhanced_array,
            precomputed_noisy_transcript,
        ],
    )

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 / "active_light.css").read_text(encoding="utf-8"),
    allowed_paths=[APP_TMP_DIR, "/tmp", "/"],
)