from nemo.collections.asr.models import ASRModel import torch import gradio as gr import spaces import gc import shutil from pathlib import Path from pydub import AudioSegment import numpy as np import os import gradio.themes as gr_themes import csv import json from typing import List, Tuple 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 start_session(request: gr.Request): session_hash = request.session_hash session_dir = Path(f'/tmp/{session_hash}') session_dir.mkdir(parents=True, exist_ok=True) print(f"Session with hash {session_hash} started.") return session_dir.as_posix() def end_session(request: gr.Request): session_hash = request.session_hash session_dir = Path(f'/tmp/{session_hash}') if session_dir.exists(): shutil.rmtree(session_dir) print(f"Session with hash {session_hash} ended.") def get_audio_segment(audio_path, start_second, end_second): 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 def preprocess_audio(audio_path, session_dir): """ オーディオファイルの前処理（リサンプリング、モノラル変換）を行う。 Args: audio_path (str): 入力オーディオファイルのパス。 session_dir (str): セッションディレクトリのパス。 Returns: tuple: (processed_path, info_path_name, duration_sec) のタプル、または None（処理に失敗した場合）。 """ try: original_path_name = Path(audio_path).name audio_name = Path(audio_path).stem try: gr.Info(f"Loading audio: {original_path_name}", duration=2) audio = AudioSegment.from_file(audio_path) duration_sec = audio.duration_seconds except Exception as load_e: gr.Error(f"Failed to load audio file {original_path_name}: {load_e}", duration=None) return None, None, None 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 None, None, None 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 None, None, None elif audio.channels > 2: gr.Error(f"Audio has {audio.channels} channels. Only mono (1) or stereo (2) supported.", duration=None) return None, None, None processed_audio_path = None if resampled or mono: try: processed_audio_path = Path(session_dir, f"{audio_name}_resampled.wav") audio.export(processed_audio_path, format="wav") transcribe_path = processed_audio_path.as_posix() 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 processed_audio_path and os.path.exists(processed_audio_path): os.remove(processed_audio_path) return None, None, None else: transcribe_path = audio_path info_path_name = original_path_name return transcribe_path, info_path_name, duration_sec except Exception as e: gr.Error(f"Audio preprocessing failed: {e}", duration=None) return None, None, None def transcribe_audio(transcribe_path, model, duration_sec, device): """ オーディオファイルを文字起こしし、タイムスタンプを取得する。 Args: transcribe_path (str): 入力オーディオファイルのパス。 model (ASRModel): 使用するASRモデル。 duration_sec (float): オーディオファイルの長さ（秒）。 device (str): 使用するデバイス（'cuda' or 'cpu'）。 Returns: tuple: (vis_data, raw_times_data, word_vis_data) のタプル、または None（処理に失敗した場合）。 """ long_audio_settings_applied = False try: # CUDA 使用前にメモリをクリアし、断片化を低減 if device == 'cuda': torch.cuda.empty_cache() gc.collect() model.to(device) model.to(torch.float32) # メモリ状況をログ出力（デバッグ用） if device == 'cuda': print(f"CUDA Memory before transcription: {torch.cuda.memory_allocated() / 1024**2:.2f} MB") gr.Info(f"Transcribing on {device}...", duration=2) if duration_sec > 480: 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) 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}") 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 'segment' not in output[0].timestamp: gr.Error("Transcription failed or produced unexpected output format.", duration=None) return None, None, None segment_timestamps = output[0].timestamp['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] word_timestamps_raw = output[0].timestamp.get("word", []) word_vis_data = [ [f"{w['start']:.2f}", f"{w['end']:.2f}", w["word"]] for w in word_timestamps_raw if isinstance(w, dict) and 'start' in w and 'end' in w and 'word' in w ] gr.Info("Transcription complete.", duration=2) return vis_data, raw_times_data, word_vis_data 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 None, None, None except Exception as e: error_msg = f"Transcription failed: {e}" print(f"Error during transcription processing: {e}") gr.Error(error_msg, duration=None) return None, None, None finally: try: if long_audio_settings_applied: try: print("Reverting long audio settings.") model.change_attention_model("rel_pos") model.change_subsampling_conv_chunking_factor(-1) 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) 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) def save_transcripts(session_dir, audio_name, vis_data, word_vis_data): """ 文字起こし結果を各種ファイル形式（CSV、SRT、VTT、JSON、LRC）で保存する。 Args: session_dir (str): セッションディレクトリのパス。 audio_name (str): オーディオファイルの名前。 vis_data (list): 表示用の文字起こし結果のリスト。 word_vis_data (list): 単語レベルのタイムスタンプのリスト。 Returns: tuple: 各ファイルのダウンロードボタンの更新情報を含むタプル。 """ try: csv_headers = ["Start (s)", "End (s)", "Segment"] csv_file_path = Path(session_dir, f"transcription_{audio_name}.csv") with open(csv_file_path, 'w', newline='', encoding='utf-8') as f: writer = csv.writer(f) writer.writerow(csv_headers) writer.writerows(vis_data) print(f"CSV transcript saved to temporary file: {csv_file_path}") srt_file_path = Path(session_dir, f"transcription_{audio_name}.srt") vtt_file_path = Path(session_dir, f"transcription_{audio_name}.vtt") json_file_path = Path(session_dir, f"transcription_{audio_name}.json") write_srt(vis_data, srt_file_path) write_vtt(vis_data, word_vis_data, vtt_file_path) write_json(vis_data, word_vis_data, json_file_path) print(f"SRT, VTT, JSON transcript saved to temporary files: {srt_file_path}, {vtt_file_path}, {json_file_path}") lrc_file_path = Path(session_dir, f"transcription_{audio_name}.lrc") write_lrc(vis_data, lrc_file_path) print(f"LRC transcript saved to temporary file: {lrc_file_path}") return ( gr.DownloadButton(value=csv_file_path.as_posix(), visible=True), gr.DownloadButton(value=srt_file_path.as_posix(), visible=True), gr.DownloadButton(value=vtt_file_path.as_posix(), visible=True), gr.DownloadButton(value=json_file_path.as_posix(), visible=True), gr.DownloadButton(value=lrc_file_path.as_posix(), visible=True) ) except Exception as e: gr.Error(f"Failed to create transcript files: {e}", duration=None) print(f"Error writing transcript files: {e}") return tuple([gr.DownloadButton(visible=False)] * 5) def split_audio_with_overlap(audio_path: str, session_dir: str, chunk_length_sec: int = 3600, overlap_sec: int = 30) -> List[str]: """ 音声ファイルをchunk_length_secごとにoverlap_secのオーバーラップ付きで分割し、 分割ファイルのパスリストを返す。 """ audio = AudioSegment.from_file(audio_path) duration = audio.duration_seconds chunk_paths = [] start = 0 chunk_idx = 0 while start < duration: end = min(start + chunk_length_sec, duration) # オーバーラップを考慮 chunk_start = max(0, start - (overlap_sec if start > 0 else 0)) chunk_end = min(end + (overlap_sec if end < duration else 0), duration) chunk = audio[chunk_start * 1000:chunk_end * 1000] chunk_path = Path(session_dir, f"chunk_{chunk_idx:03d}.wav").as_posix() chunk.export(chunk_path, format="wav") chunk_paths.append(chunk_path) start += chunk_length_sec chunk_idx += 1 return chunk_paths @spaces.GPU def get_transcripts_and_raw_times(audio_path, session_dir): """ オーディオファイルを処理し、文字起こし結果を生成する。 3時間を超える場合は60分ごとに分割し、オーバーラップ付きでASRを実行してマージする。 """ if not audio_path: gr.Error("No audio file path provided for transcription.", duration=None) return [], [], [], None, gr.DownloadButton(visible=False), gr.DownloadButton(visible=False), gr.DownloadButton(visible=False), gr.DownloadButton(visible=False), gr.DownloadButton(visible=False) audio_name = Path(audio_path).stem processed_audio_path = None temp_chunk_paths = [] try: # オーディオの前処理 transcribe_path, info_path_name, duration_sec = preprocess_audio(audio_path, session_dir) if not transcribe_path or not duration_sec: return [], [], [], audio_path, gr.DownloadButton(visible=False), gr.DownloadButton(visible=False), gr.DownloadButton(visible=False), gr.DownloadButton(visible=False), gr.DownloadButton(visible=False) processed_audio_path = transcribe_path if transcribe_path != audio_path else None # 3時間超の場合は分割して逐次ASR if duration_sec > 10800: chunk_paths = split_audio_with_overlap(transcribe_path, session_dir, chunk_length_sec=3600, overlap_sec=30) temp_chunk_paths = chunk_paths.copy() all_vis_data = [] all_raw_times_data = [] all_word_vis_data = [] offset = 0.0 prev_end = 0.0 for i, chunk_path in enumerate(chunk_paths): chunk_audio = AudioSegment.from_file(chunk_path) chunk_duration = chunk_audio.duration_seconds # ASR実行 result = transcribe_audio(chunk_path, model, chunk_duration, device) if not result: continue vis_data, raw_times_data, word_vis_data = result # タイムスタンプを全体のオフセットに合わせて補正 vis_data_offset = [] raw_times_data_offset = [] word_vis_data_offset = [] for row in vis_data: s, e, seg = float(row[0]), float(row[1]), row[2] vis_data_offset.append([f"{s+offset:.2f}", f"{e+offset:.2f}", seg]) for row in raw_times_data: s, e = float(row[0]), float(row[1]) raw_times_data_offset.append([s+offset, e+offset]) for row in word_vis_data: s, e, w = float(row[0]), float(row[1]), row[2] word_vis_data_offset.append([f"{s+offset:.2f}", f"{e+offset:.2f}", w]) # オーバーラップ部分の重複除去（単純に前回のend以降のみ追加） vis_data_offset = [row for row in vis_data_offset if float(row[0]) >= prev_end] raw_times_data_offset = [row for row in raw_times_data_offset if row[0] >= prev_end] word_vis_data_offset = [row for row in word_vis_data_offset if float(row[0]) >= prev_end] if vis_data_offset: prev_end = float(vis_data_offset[-1][1]) all_vis_data.extend(vis_data_offset) all_raw_times_data.extend(raw_times_data_offset) all_word_vis_data.extend(word_vis_data_offset) offset += chunk_duration - (30 if i < len(chunk_paths)-1 else 0) # ファイルの保存 button_updates = save_transcripts(session_dir, audio_name, all_vis_data, all_word_vis_data) # 一時分割ファイル削除 for p in temp_chunk_paths: try: os.remove(p) except Exception: pass return ( all_vis_data, all_raw_times_data, all_word_vis_data, audio_path, *button_updates ) else: # 3時間以内は従来通り result = transcribe_audio(transcribe_path, model, duration_sec, device) if not result: return [], [], [], audio_path, gr.DownloadButton(visible=False), gr.DownloadButton(visible=False), gr.DownloadButton(visible=False), gr.DownloadButton(visible=False), gr.DownloadButton(visible=False) vis_data, raw_times_data, word_vis_data = result button_updates = save_transcripts(session_dir, audio_name, vis_data, word_vis_data) return ( vis_data, raw_times_data, word_vis_data, audio_path, *button_updates ) 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}") # 分割ファイルの掃除 for p in temp_chunk_paths: if os.path.exists(p): try: os.remove(p) except Exception: pass def play_segment(evt: gr.SelectData, raw_ts_list, current_audio_path): 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") def write_srt(segments, path): def sec2srt(t): h, rem = divmod(int(float(t)), 3600) m, s = divmod(rem, 60) ms = int((float(t) - int(float(t))) * 1000) return f"{h:02}:{m:02}:{s:02},{ms:03}" with open(path, "w", encoding="utf-8") as f: for i, seg in enumerate(segments, 1): f.write(f"{i}\n{sec2srt(seg[0])} --> {sec2srt(seg[1])}\n{seg[2]}\n\n") def write_vtt(segments, words, path): def sec2vtt(t): h, rem = divmod(int(float(t)), 3600) m, s = divmod(rem, 60) ms = int((float(t) - int(float(t))) * 1000) return f"{h:02}:{m:02}:{s:02}.{ms:03}" with open(path, "w", encoding="utf-8") as f: f.write("WEBVTT\n\n") word_idx = 0 for seg in segments: s_start = float(seg[0]) s_end = float(seg[1]) s_text = seg[2] # このセグメントに含まれる単語を抽出 segment_words = [] while word_idx < len(words): w = words[word_idx] w_start = float(w[0]) w_end = float(w[1]) if w_start >= s_start and w_end <= s_end: segment_words.append(w) word_idx += 1 elif w_end < s_start: word_idx += 1 else: break # 各単語ごとにタイムスタンプを生成 for i, w in enumerate(segment_words): w_start = float(w[0]) w_end = float(w[1]) w_text = w[2] # 現在の単語を強調表示し、他の単語は通常表示 colored_text = "" for j, other_w in enumerate(segment_words): if j == i: colored_text += f"{other_w[2]} " else: colored_text += f"{other_w[2]} " f.write(f"{sec2vtt(w_start)} --> {sec2vtt(w_end)}\n{colored_text.strip()}\n\n") def write_json(segments, words, path): result = {"segments": []} word_idx = 0 for s in segments: s_start = float(s[0]) s_end = float(s[1]) s_text = s[2] word_list = [] while word_idx < len(words): w = words[word_idx] w_start = float(w[0]) w_end = float(w[1]) if w_start >= s_start and w_end <= s_end: word_list.append({"start": w_start, "end": w_end, "word": w[2]}) word_idx += 1 elif w_end < s_start: word_idx += 1 else: break result["segments"].append({ "start": s_start, "end": s_end, "text": s_text, "words": word_list }) with open(path, "w", encoding="utf-8") as f: json.dump(result, f, ensure_ascii=False, indent=2) def write_lrc(segments, path): def sec2lrc(t): m, s = divmod(float(t), 60) return f"[{int(m):02}:{s:05.2f}]" with open(path, "w", encoding="utf-8") as f: for seg in segments: f.write(f"{sec2lrc(seg[0])}{seg[2]}\n") article = ( "

" "This demo showcases parakeet-tdt-0.6b-v2, a 600M-parameter model for high-quality English ASR.
" "Now optimised for long recordings (hours) with automatic chunking & memory control." "

" "

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"], ] nvidia_theme = gr_themes.Default( primary_hue=gr_themes.Color( c50="#E6F1D9", c100="#CEE3B3", c200="#B5D58C", c300="#9CC766", c400="#84B940", c500="#76B900", c600="#68A600", c700="#5A9200", c800="#4C7E00", c900="#3E6A00", c950="#2F5600" ), neutral_hue="gray", font=[gr_themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"], ).set() 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} (Long-audio ready)

") gr.HTML(article) current_audio_path_state = gr.State(None) raw_timestamps_list_state = gr.State([]) session_dir_state = gr.State() demo.load(start_session, outputs=[session_dir_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

") download_btn = gr.DownloadButton(label="Download Segment Transcript (CSV)", visible=False) srt_btn = gr.DownloadButton(label="Download SRT", visible=False) vtt_btn = gr.DownloadButton(label="Download VTT", visible=False) json_btn = gr.DownloadButton(label="Download JSON", visible=False) lrc_btn = gr.DownloadButton(label="Download LRC", visible=False) with gr.Tabs(): with gr.TabItem("Segment View (Click row to play segment)"): vis_timestamps_df = gr.DataFrame( headers=["Start (s)", "End (s)", "Segment"], datatype=["number", "number", "str"], wrap=True, ) selected_segment_player = gr.Audio(label="Selected Segment", interactive=False) with gr.TabItem("Word View"): word_vis_df = gr.DataFrame( headers=["Start (s)", "End (s)", "Word"], datatype=["number", "number", "str"], wrap=False, ) mic_transcribe_btn.click( fn=get_transcripts_and_raw_times, inputs=[mic_input, session_dir_state], outputs=[vis_timestamps_df, raw_timestamps_list_state, word_vis_df, current_audio_path_state, download_btn, srt_btn, vtt_btn, json_btn, lrc_btn], api_name="transcribe_mic" ) file_transcribe_btn.click( fn=get_transcripts_and_raw_times, inputs=[file_input, session_dir_state], outputs=[vis_timestamps_df, raw_timestamps_list_state, word_vis_df, current_audio_path_state, download_btn, srt_btn, vtt_btn, json_btn, lrc_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], ) demo.unload(end_session) if __name__ == "__main__": print("Launching Gradio Demo...") # タイムアウト対策としてキューサイズと同時実行数を抑制 demo.queue( max_size=5, default_concurrency_limit=1 ) demo.launch()