import os, whisperx, gc

device       = "cpu"          # ★CPU 固定
audio_file   = "sample.wav"   # 数十秒の WAV/16 kHz/mono
batch_size   = 4              # CPU なので小さめ
compute_type = "int8"         # メモリ最小

# 1) ASR
model = whisperx.load_model("medium", device, compute_type=compute_type,
                            language="ja")               # ←自動判定オフで高速
audio = whisperx.load_audio(audio_file)
result = model.transcribe(audio, batch_size=batch_size)

# 2) 強制アライン
align_model, meta = whisperx.load_align_model(language_code="ja", device=device)
result = whisperx.align(result["segments"], align_model, meta,
                        audio, device, return_char_alignments=False)
del align_model; gc.collect()

# 3) ダイアリゼーション
hf_token = os.getenv("HF_TOKEN")        # 事前に export HF_TOKEN=... しておく
dia = whisperx.diarize.DiarizationPipeline(
        use_auth_token=hf_token,
        device=device)
dia_segments = dia(audio_file)          # ← WAV パスを渡す
result = whisperx.assign_word_speakers(dia_segments, result)

# 4) 出力
for seg in result["segments"]:
    print(f"{seg['start']:.2f}-{seg['end']:.2f} {seg.get('speaker')}: {seg['text']}")