# NeMo ASRモデル、PyTorch、Gradioなどをインポート from nemo.collections.asr.models import ASRModel import torch import gradio as gr import spaces # Hugging Face Spaces ライブラリをインポート import gc from pathlib import Path import os import json from typing import List, Tuple, Optional # pydub をインポート (音声ファイルの長さ取得のため) try: from pydub import AudioSegment PYDUB_AVAILABLE = True except ImportError: PYDUB_AVAILABLE = False print("Warning: pydub not found. Audio duration cannot be determined automatically for long audio optimization.") # グローバル設定 MODEL_NAME = "nvidia/parakeet-tdt-0.6b-v2" TARGET_SAMPLE_RATE = 16000 # 音声の長さに関する閾値 (秒) LONG_AUDIO_THRESHOLD_SECONDS = 480 # 8分 VERY_LONG_AUDIO_THRESHOLD_SECONDS = 10800 # 3時間 # チャンク分割時の設定 CHUNK_LENGTH_SECONDS = 1800 # 30分 CHUNK_OVERLAP_SECONDS = 60 # 1分 # セグメント処理の設定 MAX_SEGMENT_LENGTH_SECONDS = 15 # 最大セグメント長（秒）を15秒に短縮 MAX_SEGMENT_CHARS = 100 # 最大セグメント文字数を100文字に短縮 MIN_SEGMENT_GAP_SECONDS = 0.3 # 最小セグメント間隔（秒） # VTTファイルの最大サイズ（バイト） MAX_VTT_SIZE_BYTES = 10 * 1024 * 1024 # 10MB # 文の区切り文字 SENTENCE_ENDINGS = ['.', '!', '?', '。', '！', '？'] SENTENCE_PAUSES = [',', '、', ';', '；', ':', '：'] device = "cuda" if torch.cuda.is_available() else "cpu" # スクリプト起動時のデバイス検出 # モデルの初期化 (グローバルに一度だけ行う) print(f"Initializing ASR model: {MODEL_NAME}") print(f"Initial device check: {device}") # 初期デバイス確認 model = ASRModel.from_pretrained(model_name=MODEL_NAME) model.eval() # 初期状態ではモデルをCPUに置いておく (GPU関数内で .to(device) する) model.cpu() print("ASR model initialized and moved to CPU.") def find_natural_break_point(text: str, max_length: int) -> int: """テキスト内で自然な区切り点を探す""" if len(text) <= max_length: return len(text) # 文末で区切る for i in range(max_length, 0, -1): if i < len(text) and text[i] in SENTENCE_ENDINGS: return i + 1 # 文の区切りで区切る for i in range(max_length, 0, -1): if i < len(text) and text[i] in SENTENCE_PAUSES: return i + 1 # スペースで区切る for i in range(max_length, 0, -1): if i < len(text) and text[i].isspace(): return i + 1 # それでも見つからない場合は最大長で区切る return max_length def split_segment(segment: dict, max_length_seconds: float, max_chars: int) -> List[dict]: """セグメントを自然な区切りで分割する""" if (segment['end'] - segment['start']) <= max_length_seconds and len(segment['segment']) <= max_chars: return [segment] result = [] current_text = segment['segment'] current_start = segment['start'] total_duration = segment['end'] - segment['start'] while current_text: # 文字数に基づく分割点を探す break_point = find_natural_break_point(current_text, max_chars) # 時間に基づく分割点を計算 text_ratio = break_point / len(segment['segment']) segment_duration = total_duration * text_ratio # 分割点が最大長を超えないように調整 if segment_duration > max_length_seconds: time_ratio = max_length_seconds / total_duration break_point = int(len(segment['segment']) * time_ratio) break_point = find_natural_break_point(current_text, break_point) segment_duration = max_length_seconds # 新しいセグメントを作成 new_segment = { 'start': current_start, 'end': current_start + segment_duration, 'segment': current_text[:break_point].strip() } result.append(new_segment) # 残りのテキストと開始時間を更新 current_text = current_text[break_point:].strip() current_start = new_segment['end'] return result def transcribe_audio_core( audio_path: str, duration_sec: float, current_device: str # 実行時のデバイスを引数で受け取る ) -> Tuple[Optional[List], Optional[List], Optional[List]]: """ 音声ファイルを文字起こしし、タイムスタンプを取得する（コア処理）。 この関数は実際にGPU上で実行されることを想定。 """ long_audio_settings_applied = False try: gr.Info(f"Starting transcription on {current_device} for: {Path(audio_path).name}", duration=3) if current_device == 'cuda': torch.cuda.empty_cache() gc.collect() print(f"CUDA memory before loading model to GPU: {torch.cuda.memory_allocated() / 1024**2:.2f} MB") # モデルを実行デバイスに移動 model.to(current_device) model.to(torch.float32) # 推論前にfloat32に戻す (bfloat16は後段) if current_device == 'cuda': print(f"CUDA memory after loading model to GPU (float32): {torch.cuda.memory_allocated() / 1024**2:.2f} MB") # 長尺音声用の設定 (閾値を超え、かつpydubで長さが取得できた場合) if PYDUB_AVAILABLE and duration_sec > LONG_AUDIO_THRESHOLD_SECONDS: gr.Info(f"Audio duration ({duration_sec:.2f}s) exceeds threshold. Applying long audio settings.", duration=3) try: print("Applying long audio settings: Local Attention and Chunking.") model.change_attention_model("rel_pos_local_attn", [128, 128]) # 256,256から128,128に変更 model.change_subsampling_conv_chunking_factor(1) long_audio_settings_applied = True print("Successfully applied long audio settings.") except Exception as setting_e: warning_msg = f"Warning: Failed to apply long audio settings: {setting_e}" print(warning_msg) gr.Warning(warning_msg, duration=5) # bfloat16への変換 (CUDAの場合のみ) if current_device == 'cuda': print("Converting model to bfloat16 for inference on CUDA.") model.to(torch.bfloat16) print(f"CUDA memory after converting to bfloat16: {torch.cuda.memory_allocated() / 1024**2:.2f} MB") # 文字起こし実行 print(f"Transcribing {audio_path}...") output = model.transcribe([audio_path], timestamps=True, batch_size=2) # バッチサイズを2に設定 print("Transcription API call finished.") 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: error_msg = "Transcription failed or produced unexpected output format." print(error_msg) gr.Error(error_msg, duration=5) return None, None, None segment_timestamps = output[0].timestamp['segment'] # セグメントの前処理：より適切なセグメント分割 processed_segments = [] current_segment = None for ts in segment_timestamps: if current_segment is None: current_segment = ts else: # セグメント結合の条件を厳格化 time_gap = ts['start'] - current_segment['end'] current_text = current_segment['segment'] next_text = ts['segment'] # 結合条件のチェック should_merge = ( time_gap < MIN_SEGMENT_GAP_SECONDS and # 時間間隔が短い len(current_text) + len(next_text) < MAX_SEGMENT_CHARS and # 文字数制限 (current_segment['end'] - current_segment['start']) < MAX_SEGMENT_LENGTH_SECONDS and # 現在のセグメントが短い (ts['end'] - ts['start']) < MAX_SEGMENT_LENGTH_SECONDS and # 次のセグメントが短い not any(current_text.strip().endswith(p) for p in SENTENCE_ENDINGS) # 文の区切りでない ) if should_merge: current_segment['end'] = ts['end'] current_segment['segment'] += ' ' + ts['segment'] else: # 現在のセグメントを分割 split_segments = split_segment(current_segment, MAX_SEGMENT_LENGTH_SECONDS, MAX_SEGMENT_CHARS) processed_segments.extend(split_segments) current_segment = ts if current_segment is not None: # 最後のセグメントも分割 split_segments = split_segment(current_segment, MAX_SEGMENT_LENGTH_SECONDS, MAX_SEGMENT_CHARS) processed_segments.extend(split_segments) # 処理済みセグメントからデータを生成 vis_data = [[f"{ts['start']:.2f}", f"{ts['end']:.2f}", ts['segment']] for ts in processed_segments] raw_times_data = [[ts['start'], ts['end']] for ts in processed_segments] # 単語タイムスタンプの処理を改善 word_timestamps_raw = output[0].timestamp.get("word", []) word_vis_data = [] for w in word_timestamps_raw: if not isinstance(w, dict) or not all(k in w for k in ['start', 'end', 'word']): continue # 単語のタイムスタンプを最も近いセグメントに割り当て word_start = float(w['start']) word_end = float(w['end']) # 単語が完全に含まれるセグメントを探す for seg in processed_segments: if word_start >= seg['start'] - 0.05 and word_end <= seg['end'] + 0.05: word_vis_data.append([f"{word_start:.2f}", f"{word_end:.2f}", w["word"]]) break gr.Info("Transcription successful!", duration=3) return vis_data, raw_times_data, word_vis_data except torch.cuda.OutOfMemoryError as oom_e: error_msg = f"CUDA out of memory during transcription: {oom_e}. Try a shorter audio file or a more powerful GPU." print(error_msg) gr.Error(error_msg, duration=None) # 長く表示 return None, None, None except Exception as e: error_msg = f"Error during transcription: {e}" print(error_msg) gr.Error(error_msg, duration=None) # 長く表示 return None, None, None finally: print("Starting transcription cleanup...") if long_audio_settings_applied: try: print("Reverting long audio settings...") model.change_attention_model("rel_pos") # 元のAttentionに戻す model.change_subsampling_conv_chunking_factor(-1) # 元のChunking Factorに戻す print("Successfully reverted long audio settings.") except Exception as revert_e: warning_msg = f"Warning: Failed to revert long audio settings: {revert_e}" print(warning_msg) gr.Warning(warning_msg, duration=5) # モデルをCPUに戻し、CUDAキャッシュをクリア model.cpu() # 必ずCPUに戻す print("Model moved to CPU.") if current_device == 'cuda': # current_device を使う gc.collect() torch.cuda.empty_cache() print("CUDA cache cleared.") print("Transcription cleanup finished.") @spaces.GPU(duration=60) # GPUリソースを要求し、タイムアウトを60秒に設定 def process_audio_file(audio_filepath: str) -> dict: # Gradioから渡されるのは一時ファイルのパス """ アップロードされた音声ファイルを処理し、文字起こし結果をJSONで返す。 この関数がGradioのコールバックとなり、GPU環境で実行される。 """ # この関数が呼ばれた時点でGPUが利用可能になっているはず (Hugging Face Spacesの場合) # なので、再度デバイスチェックを行う current_processing_device = "cuda" if torch.cuda.is_available() else "cpu" print(f"Device check inside @spaces.GPU function: {current_processing_device}") gr.Info(f"Processing on: {current_processing_device}", duration=3) if not PYDUB_AVAILABLE: gr.Warning("pydub library is not available. Audio duration cannot be determined, long audio optimizations might not be applied correctly.", duration=5) duration_sec = 0 # 長さが不明な場合は0とする else: try: gr.Info(f"Loading audio file: {Path(audio_filepath).name}", duration=2) audio = AudioSegment.from_file(audio_filepath) duration_sec = audio.duration_seconds print(f"Audio duration: {duration_sec:.2f} seconds.") except Exception as e: error_msg = f"Failed to load audio or get duration using pydub: {e}" print(error_msg) gr.Error(error_msg, duration=5) # pydubが失敗しても、NeMoは処理を試みることができるので、duration_sec = 0 で続行 duration_sec = 0 # 文字起こしコア処理を呼び出し vis_data, raw_times_data, word_vis_data = transcribe_audio_core(audio_filepath, duration_sec, current_processing_device) if not vis_data: # transcribe_audio_core内でエラー通知はされているはず return {"error": "Transcription failed. Check logs and messages for details."} # 結果をJSON形式で返却 (ユーザー指定の形式に合わせる) output_segments = [] word_idx = 0 for seg_data in vis_data: s_start_time = float(seg_data[0]) s_end_time = float(seg_data[1]) s_text = seg_data[2] segment_words_list: List[dict] = [] if word_vis_data: # word_vis_data が存在する場合のみ処理 temp_current_word_idx = word_idx while temp_current_word_idx < len(word_vis_data): w_data = word_vis_data[temp_current_word_idx] w_start_time = float(w_data[0]) w_end_time = float(w_data[1]) # 単語がセグメントの範囲内にあるかチェック (多少の誤差を許容) if w_start_time >= s_start_time and w_end_time <= s_end_time + 0.1: segment_words_list.append({ "start": w_start_time, "end": w_end_time, "word": w_data[2] }) temp_current_word_idx += 1 elif w_start_time < s_start_time: # 単語がセグメントより前に開始している場合はスキップ temp_current_word_idx += 1 elif w_start_time > s_end_time: # 単語がセグメントより後に開始している場合はループを抜ける break else: # その他のケース (ほぼありえないが念のため) temp_current_word_idx += 1 word_idx = temp_current_word_idx output_segments.append({ "start": s_start_time, "end": s_end_time, "text": s_text, "words": segment_words_list }) result = {"segments": output_segments} return result # Gradioインターフェースの設定 with gr.Blocks() as demo: gr.Markdown("# GPU Transcription Service (Improved)") gr.Markdown("Upload an audio file for transcription. Processing will use GPU if available on the server.") file_input = gr.File(label="Upload Audio File", type="filepath") # type="filepath" を明示 output_json = gr.JSON(label="Transcription Result") file_input.change( # ファイルがアップロード/変更されたら実行 fn=process_audio_file, inputs=[file_input], outputs=[output_json] ) gr.Examples( examples=[ [os.path.join(os.path.dirname(__file__), "audio_example.wav") if os.path.exists(os.path.join(os.path.dirname(__file__), "audio_example.wav")) else "https://www.kozco.com/tech/piano2-CoolEdit.mp3"] ], inputs=[file_input], label="Example Audio (Click to load)" ) if __name__ == "__main__": # ダミーの音声ファイルを作成 (Examples用、もし存在しなければ) example_dir = os.path.dirname(__file__) dummy_audio_path = os.path.join(example_dir, "audio_example.wav") if not os.path.exists(dummy_audio_path) and PYDUB_AVAILABLE: try: print(f"Creating a dummy audio file for example: {dummy_audio_path}") silence = AudioSegment.silent(duration=1000) # 1秒の無音 # 簡単な音を追加 (pydubの機能で) tone1 = AudioSegment.sine(440, duration=200) # A4 tone2 = AudioSegment.sine(880, duration=200) # A5 dummy_segment = silence + tone1 + silence[:200] + tone2 + silence dummy_segment.export(dummy_audio_path, format="wav") print("Dummy audio file created.") except Exception as e: print(f"Could not create dummy audio file: {e}") elif not PYDUB_AVAILABLE: print("Skipping dummy audio file creation as pydub is not available.") print("Launching Gradio demo...") demo.queue() # リクエストキューを有効化 demo.launch(show_error=True) # エラー詳細を表示