from nemo.collections.asr.models import ASRModel import torch import gradio as gr import spaces import gc import tempfile from pathlib import Path import os import subprocess import gradio.themes as gr_themes import time device = "cuda" if torch.cuda.is_available() else "cpu" EXPORT_ROOT = Path("/tmp/parakeet_api") EXPORT_ROOT.mkdir(parents=True, exist_ok=True) model = ASRModel.from_pretrained(model_name="nvidia/parakeet-tdt-0.6b-v3") model.eval() def create_request_dir() -> Path: return Path(tempfile.mkdtemp(prefix="parakeet_", dir=EXPORT_ROOT.as_posix())) def get_audio_duration_seconds(audio_path: str) -> float: ffprobe_cmd = [ "ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "default=noprint_wrappers=1:nokey=1", audio_path, ] try: result = subprocess.run(ffprobe_cmd, check=True, capture_output=True, text=True) return float(result.stdout.strip()) except Exception as e: raise gr.Error(f"Failed to inspect audio duration: {e}") from e def prepare_audio_for_transcription(audio_path: str, request_dir: Path): if not audio_path: raise gr.Error("No audio file path provided for transcription.") original_path_name = Path(audio_path).name audio_name = Path(audio_path).stem processed_audio_path = request_dir / f"{audio_name}_prepared.wav" ffmpeg_cmd = [ "ffmpeg", "-y", "-i", audio_path, "-vn", "-ac", "1", "-ar", "16000", processed_audio_path.as_posix(), ] try: subprocess.run(ffmpeg_cmd, check=True, capture_output=True, text=True) except Exception as e: raise gr.Error(f"FFmpeg preprocessing failed: {e}") from e duration_sec = get_audio_duration_seconds(processed_audio_path.as_posix()) use_long_audio_settings = duration_sec > 480 return processed_audio_path, duration_sec, f"{original_path_name} (processed)", use_long_audio_settings def extract_timestamp_text(ts: dict) -> str: return ts.get("word") or ts.get("segment") or ts.get("char") or "" def build_transcription_response(words=None): return { "source": "nvidia/parakeet-tdt-0.6b-v3", "words": words or [], } @spaces.GPU def run_transcription_on_preprocessed_audio(transcribe_path, duration_sec, info_path_name, use_long_audio_settings): # Flag to track if long audio settings were applied long_audio_settings_applied = False try: model.to(device) model.to(torch.float32) gr.Info(f"Transcribing {info_path_name} on {device}...", duration=2) # Apply long audio settings only when preprocessing determined they are needed if use_long_audio_settings: try: gr.Info("Audio longer than 8 minutes. Applying optimized settings for long transcription.", duration=3) print("Applying long audio settings: Local Attention and Chunking.") model.change_attention_model("rel_pos_local_attn", [256,256]) model.change_subsampling_conv_chunking_factor(1) # 1 = auto select long_audio_settings_applied = True except Exception as setting_e: gr.Warning(f"Could not apply long audio settings: {setting_e}", duration=5) print(f"Warning: Failed to apply long audio settings: {setting_e}") # Proceed without long audio settings if applying them failed model.to(torch.bfloat16) 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 'word' not in output[0].timestamp: raise gr.Error("Word Timestamp Format Issue") word_timestamps = output[0].timestamp['word'] words = [] for idx, ts in enumerate(word_timestamps, start=1): word_text = extract_timestamp_text(ts) start = round(float(ts["start"]), 3) end = round(float(ts["end"]), 3) words.append({ "word_id": f"word_{idx:06d}", "text": word_text, "start": start, "end": end, "duration": round(end - start, 3), }) gr.Info("Transcription complete.", duration=2) return build_transcription_response(words=words) 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}") raise gr.Error(error_msg) from e 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}") raise gr.Error("File not found for transcription") except Exception as e: error_msg = f"Transcription failed: {e}" print(f"Error during transcription processing: {e}") raise gr.Error(error_msg) from e finally: # --- Model Cleanup --- try: # Revert settings if they were applied for long audio if long_audio_settings_applied: try: print("Reverting long audio settings.") model.change_attention_model("rel_pos") model.change_subsampling_conv_chunking_factor(-1) long_audio_settings_applied = False # Reset flag except Exception as revert_e: print(f"Warning: Failed to revert long audio settings: {revert_e}") gr.Warning(f"Issue reverting model settings after long transcription: {revert_e}", duration=5) gc.collect() except Exception as cleanup_e: print(f"Error during model cleanup: {cleanup_e}") gr.Warning(f"Issue during model cleanup: {cleanup_e}", duration=5) # --- End Model Cleanup --- def get_transcripts_and_raw_times(audio_path): request_dir = create_request_dir() processed_audio_path = None try: processed_audio_path, duration_sec, info_path_name, use_long_audio_settings = prepare_audio_for_transcription(audio_path, request_dir) started_at = time.perf_counter() result = run_transcription_on_preprocessed_audio( processed_audio_path.as_posix(), duration_sec, info_path_name, use_long_audio_settings, ) result["zerogpu_seconds"] = round(time.perf_counter() - started_at, 3) return result 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}") article = ( "

" "This API-first demo uses parakeet-tdt-0.6b-v3 to transcribe uploaded audio files with word-level timestamps." "

" "

" "Upload an audio file to get word-level timestamps in a structured JSON response. This Space is optimized for file uploads and API usage only; microphone/live transcription, transcript downloads, and preview playback have been removed." "

" "

" "🎙️ Model Card | " "🧑‍💻 NeMo Repository" "

" ) examples = [ ["data/example-yt_saTD1u8PorI.mp3"], ] # Define an EU-inspired theme nvidia_theme = gr_themes.Default( primary_hue=gr_themes.Color( c50="#E6ECF7", c100="#CCD9EF", c200="#99B3DF", c300="#668DCC", c400="#3366B3", c500="#003399", # EU Blue c600="#002E8A", c700="#00246D", c800="#001A51", c900="#001238", c950="#000B24" ), neutral_hue="gray", 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 = "parakeet-tdt-0.6b-v3" gr.Markdown(f"

Speech Transcription with {model_display_name} 🦜

") gr.HTML(article) 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") gr.Markdown("---") gr.Markdown("

Structured transcription response

") transcription_response_output = gr.JSON(label="Transcription Response") file_transcribe_btn.click( fn=get_transcripts_and_raw_times, inputs=[file_input], outputs=[transcription_response_output], api_name="transcribe_file" ) if __name__ == "__main__": print("Launching Gradio Demo...") demo.queue() demo.launch()