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( '' '' "" ) 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", "/"], )