from nemo.collections.asr.models import ASRModel import torch import gradio as gr import spaces import gc from pathlib import Path from pydub import AudioSegment import numpy as np import os import tempfile import gradio.themes as gr_themes import csv device = "cuda" if torch.cuda.is_available() else "cpu" MODEL_NAME="nvidia/parakeet-tdt-0.6b-v2" model = ASRModel.from_pretrained(model_name=MODEL_NAME) model.eval() def get_audio_segment(audio_path, start_second, end_second): """ Extract a segment of audio from a given audio file. Parameters: audio_path (str): Path to the audio file to process start_second (float): Start time of the segment in seconds end_second (float): End time of the segment in seconds Returns: tuple or None: A tuple containing (frame_rate, samples) where: - frame_rate (int): The sample rate of the audio - samples (numpy.ndarray): The audio samples as a numpy array Returns None if there's an error processing the audio """ if not audio_path or not Path(audio_path).exists(): print(f"Warning: Audio path '{audio_path}' not found or invalid for clipping.") return None try: start_ms = int(start_second * 1000) end_ms = int(end_second * 1000) start_ms = max(0, start_ms) if end_ms <= start_ms: print(f"Warning: End time ({end_second}s) is not after start time ({start_second}s). Adjusting end time.") end_ms = start_ms + 100 audio = AudioSegment.from_file(audio_path) clipped_audio = audio[start_ms:end_ms] samples = np.array(clipped_audio.get_array_of_samples()) if clipped_audio.channels == 2: samples = samples.reshape((-1, 2)).mean(axis=1).astype(samples.dtype) frame_rate = clipped_audio.frame_rate if frame_rate <= 0: print(f"Warning: Invalid frame rate ({frame_rate}) detected for clipped audio.") frame_rate = audio.frame_rate if samples.size == 0: print(f"Warning: Clipped audio resulted in empty samples array ({start_second}s to {end_second}s).") return None return (frame_rate, samples) except FileNotFoundError: print(f"Error: Audio file not found at path: {audio_path}") return None except Exception as e: print(f"Error clipping audio {audio_path} from {start_second}s to {end_second}s: {e}") return None @spaces.GPU def get_transcripts_and_raw_times(audio_path): """ Transcribe an audio file and generate timestamps for each segment. Parameters: audio_path (str): Path to the audio file to transcribe Returns: tuple: A tuple containing: - vis_data (list): List of [start, end, text] for visualization - raw_times_data (list): List of [start, end] timestamps - audio_path (str): Path to the processed audio file - button_update (gr.DownloadButton): Gradio button component for CSV download Notes: - Automatically handles audio preprocessing (resampling to 16kHz, mono conversion) - Uses NVIDIA's Parakeet TDT model for transcription - Generates a CSV file with transcription results """ if not audio_path: gr.Error("No audio file path provided for transcription.", duration=None) # Return an update to hide the button return [], [], None, gr.DownloadButton(visible=False) vis_data = [["N/A", "N/A", "Processing failed"]] raw_times_data = [[0.0, 0.0]] processed_audio_path = None temp_file = None csv_file_path = None original_path_name = Path(audio_path).name try: try: gr.Info(f"Loading audio: {original_path_name}", duration=2) audio = AudioSegment.from_file(audio_path) except Exception as load_e: gr.Error(f"Failed to load audio file {original_path_name}: {load_e}", duration=None) # Return an update to hide the button return [["Error", "Error", "Load failed"]], [[0.0, 0.0]], audio_path, gr.DownloadButton(visible=False) resampled = False mono = False target_sr = 16000 if audio.frame_rate != target_sr: try: audio = audio.set_frame_rate(target_sr) resampled = True except Exception as resample_e: gr.Error(f"Failed to resample audio: {resample_e}", duration=None) # Return an update to hide the button return [["Error", "Error", "Resample failed"]], [[0.0, 0.0]], audio_path, gr.DownloadButton(visible=False) if audio.channels == 2: try: audio = audio.set_channels(1) mono = True except Exception as mono_e: gr.Error(f"Failed to convert audio to mono: {mono_e}", duration=None) # Return an update to hide the button return [["Error", "Error", "Mono conversion failed"]], [[0.0, 0.0]], audio_path, gr.DownloadButton(visible=False) elif audio.channels > 2: gr.Error(f"Audio has {audio.channels} channels. Only mono (1) or stereo (2) supported.", duration=None) # Return an update to hide the button return [["Error", "Error", f"{audio.channels}-channel audio not supported"]], [[0.0, 0.0]], audio_path, gr.DownloadButton(visible=False) if resampled or mono: try: temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav") audio.export(temp_file.name, format="wav") processed_audio_path = temp_file.name temp_file.close() transcribe_path = processed_audio_path info_path_name = f"{original_path_name} (processed)" except Exception as export_e: gr.Error(f"Failed to export processed audio: {export_e}", duration=None) if temp_file and hasattr(temp_file, 'name') and os.path.exists(temp_file.name): # Check temp_file has 'name' attribute os.remove(temp_file.name) # Return an update to hide the button return [["Error", "Error", "Export failed"]], [[0.0, 0.0]], audio_path, gr.DownloadButton(visible=False) else: transcribe_path = audio_path info_path_name = original_path_name try: model.to(device) gr.Info(f"Transcribing {info_path_name} on {device}...", duration=2) output = model.transcribe([transcribe_path], timestamps=True) if not output or not isinstance(output, list) or not output[0] or not hasattr(output[0], 'timestamp') or not output[0].timestamp or 'segment' not in output[0].timestamp: gr.Error("Transcription failed or produced unexpected output format.", duration=None) # Return an update to hide the button return [["Error", "Error", "Transcription Format Issue"]], [[0.0, 0.0]], audio_path, gr.DownloadButton(visible=False) segment_timestamps = output[0].timestamp['segment'] csv_headers = ["Start (s)", "End (s)", "Segment"] vis_data = [[f"{ts['start']:.2f}", f"{ts['end']:.2f}", ts['segment']] for ts in segment_timestamps] raw_times_data = [[ts['start'], ts['end']] for ts in segment_timestamps] # Default button update (hidden) in case CSV writing fails button_update = gr.DownloadButton(visible=False) try: temp_csv_file = tempfile.NamedTemporaryFile(delete=False, suffix=".csv", mode='w', newline='', encoding='utf-8') writer = csv.writer(temp_csv_file) writer.writerow(csv_headers) writer.writerows(vis_data) csv_file_path = temp_csv_file.name temp_csv_file.close() print(f"CSV transcript saved to temporary file: {csv_file_path}") # If CSV is saved, create update to show button with path button_update = gr.DownloadButton(value=csv_file_path, visible=True) except Exception as csv_e: gr.Error(f"Failed to create transcript CSV file: {csv_e}", duration=None) print(f"Error writing CSV: {csv_e}") # csv_file_path remains None, button_update remains hidden gr.Info("Transcription complete.", duration=2) # Return the data and the button update dictionary return vis_data, raw_times_data, audio_path, button_update except torch.cuda.OutOfMemoryError as e: error_msg = 'CUDA out of memory. Please try a shorter audio or reduce GPU load.' print(f"CUDA OutOfMemoryError: {e}") gr.Error(error_msg, duration=None) # Return an update to hide the button return [["OOM", "OOM", error_msg]], [[0.0, 0.0]], audio_path, gr.DownloadButton(visible=False) except FileNotFoundError: error_msg = f"Audio file for transcription not found: {Path(transcribe_path).name}." print(f"Error: Transcribe audio file not found at path: {transcribe_path}") gr.Error(error_msg, duration=None) # Return an update to hide the button return [["Error", "Error", "File not found for transcription"]], [[0.0, 0.0]], audio_path, gr.DownloadButton(visible=False) except Exception as e: error_msg = f"Transcription failed: {e}" print(f"Error during transcription processing: {e}") gr.Error(error_msg, duration=None) vis_data = [["Error", "Error", error_msg]] raw_times_data = [[0.0, 0.0]] # Return an update to hide the button return vis_data, raw_times_data, audio_path, gr.DownloadButton(visible=False) finally: try: if 'model' in locals() and hasattr(model, 'cpu'): if device == 'cuda': model.cpu() gc.collect() if device == 'cuda': torch.cuda.empty_cache() except Exception as cleanup_e: print(f"Error during model cleanup: {cleanup_e}") gr.Warning(f"Issue during model cleanup: {cleanup_e}", duration=5) finally: if processed_audio_path and os.path.exists(processed_audio_path): try: os.remove(processed_audio_path) print(f"Temporary audio file {processed_audio_path} removed.") except Exception as e: print(f"Error removing temporary audio file {processed_audio_path}: {e}") def play_segment(evt: gr.SelectData, raw_ts_list, current_audio_path): """ Play a selected segment from the transcription results. Parameters: evt (gr.SelectData): Gradio select event containing the index of selected segment raw_ts_list (list): List of [start, end] timestamps for all segments current_audio_path (str): Path to the current audio file being processed Returns: gr.Audio: Gradio Audio component containing the selected segment for playback Notes: - Extracts and plays the audio segment corresponding to the selected transcription - Returns None if segment extraction fails or inputs are invalid """ if not isinstance(raw_ts_list, list): print(f"Warning: raw_ts_list is not a list ({type(raw_ts_list)}). Cannot play segment.") return gr.Audio(value=None, label="Selected Segment") if not current_audio_path: print("No audio path available to play segment from.") return gr.Audio(value=None, label="Selected Segment") selected_index = evt.index[0] if selected_index < 0 or selected_index >= len(raw_ts_list): print(f"Invalid index {selected_index} selected for list of length {len(raw_ts_list)}.") return gr.Audio(value=None, label="Selected Segment") if not isinstance(raw_ts_list[selected_index], (list, tuple)) or len(raw_ts_list[selected_index]) != 2: print(f"Warning: Data at index {selected_index} is not in the expected format [start, end].") return gr.Audio(value=None, label="Selected Segment") start_time_s, end_time_s = raw_ts_list[selected_index] print(f"Attempting to play segment: {current_audio_path} from {start_time_s:.2f}s to {end_time_s:.2f}s") segment_data = get_audio_segment(current_audio_path, start_time_s, end_time_s) if segment_data: print("Segment data retrieved successfully.") return gr.Audio(value=segment_data, autoplay=True, label=f"Segment: {start_time_s:.2f}s - {end_time_s:.2f}s", interactive=False) else: print("Failed to get audio segment data.") return gr.Audio(value=None, label="Selected Segment") article = ( "

" "This demo showcases parakeet-tdt-0.6b-v2, a 600-million-parameter model designed for high-quality English speech recognition." "

" "

Key Features:

" "" "

" "This model is available for commercial and non-commercial use." "

" "

" "🎙️ Learn more about the Model | " "📄 Fast Conformer paper | " "📚 TDT paper | " "🧑‍💻 NeMo Repository" "

" ) examples = [ ["data/example-yt_saTD1u8PorI.mp3"], ] # Define an NVIDIA-inspired theme nvidia_theme = gr_themes.Default( primary_hue=gr_themes.Color( c50="#E6F1D9", # Lightest green c100="#CEE3B3", c200="#B5D58C", c300="#9CC766", c400="#84B940", c500="#76B900", # NVIDIA Green c600="#68A600", c700="#5A9200", c800="#4C7E00", c900="#3E6A00", # Darkest green c950="#2F5600" ), neutral_hue="gray", # Use gray for neutral elements font=[gr_themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"], ).set() # Apply the custom theme with gr.Blocks(theme=nvidia_theme) as demo: model_display_name = MODEL_NAME.split('/')[-1] if '/' in MODEL_NAME else MODEL_NAME gr.Markdown(f"

Speech Transcription with {model_display_name}

") gr.HTML(article) current_audio_path_state = gr.State(None) raw_timestamps_list_state = gr.State([]) with gr.Tabs(): with gr.TabItem("Audio File"): file_input = gr.Audio(sources=["upload"], type="filepath", label="Upload Audio File") gr.Examples(examples=examples, inputs=[file_input], label="Example Audio Files (Click to Load)") file_transcribe_btn = gr.Button("Transcribe Uploaded File", variant="primary") with gr.TabItem("Microphone"): mic_input = gr.Audio(sources=["microphone"], type="filepath", label="Record Audio") mic_transcribe_btn = gr.Button("Transcribe Microphone Input", variant="primary") gr.Markdown("---") gr.Markdown("

Transcription Results (Click row to play segment)

") # Define the DownloadButton *before* the DataFrame download_btn = gr.DownloadButton(label="Download Transcript (CSV)", visible=False) vis_timestamps_df = gr.DataFrame( headers=["Start (s)", "End (s)", "Segment"], datatype=["number", "number", "str"], wrap=True, label="Transcription Segments" ) # selected_segment_player was defined after download_btn previously, keep it after df for layout selected_segment_player = gr.Audio(label="Selected Segment", interactive=False) mic_transcribe_btn.click( fn=get_transcripts_and_raw_times, inputs=[mic_input], outputs=[vis_timestamps_df, raw_timestamps_list_state, current_audio_path_state, download_btn], api_name="transcribe_mic" ) file_transcribe_btn.click( fn=get_transcripts_and_raw_times, inputs=[file_input], outputs=[vis_timestamps_df, raw_timestamps_list_state, current_audio_path_state, download_btn], api_name="transcribe_file" ) vis_timestamps_df.select( fn=play_segment, inputs=[raw_timestamps_list_state, current_audio_path_state], outputs=[selected_segment_player], ) if __name__ == "__main__": print("Launching Gradio Demo...") demo.queue() demo.launch(mcp_server=True)