app.py

import torch
import gradio as gr
from transformers import pipeline
from pyannote.audio import Pipeline
import datetime
from pydub import AudioSegment
import os

# โมเดลถอดเสียง
MODEL_NAME = "biodatlab/whisper-th-medium-combined"
device = 0 if torch.cuda.is_available() else "cpu"

pipe = pipeline(
    task="automatic-speech-recognition",
    model=MODEL_NAME,
    chunk_length_s=30,
    device=device,
)

# โมเดลแยกเสียงคนพูด
HF_TOKEN = os.environ.get("HF_TOKEN")
if not HF_TOKEN:
    raise ValueError("กรุณาตั้งค่า HF_TOKEN ใน environment variables")

diarization = Pipeline.from_pretrained(
    "pyannote/speaker-diarization@2.1",
    use_auth_token=HF_TOKEN
)

def format_timestamp(seconds):
    return str(datetime.timedelta(seconds=seconds)).split(".")[0]

def process_audio(audio_file):
    if audio_file is None:
        return "กรุณาอัพโหลดไฟล์เสียง"
    
    try:
        # แยกเสียงคนพูด
        diarization_result = diarization(audio_file)
        
        # ถอดความและจัดรูปแบบผลลัพธ์
        transcription = pipe(
            audio_file, 
            generate_kwargs={"language": "<|th|>", "task": "transcribe"}, 
            batch_size=16
        )["text"]
        
        # สร้างผลลัพธ์ที่มีการแยกคนพูด
        result = []
        for turn, _, speaker in diarization_result.itertracks(yield_label=True):
            start_time = format_timestamp(turn.start)
            end_time = format_timestamp(turn.end)
            result.append(f"[{start_time} - {end_time}] {speaker}:")
            
        # รวมผลลัพธ์
        final_result = "\n".join(result) + "\n\nคำถอดความ:\n" + transcription
        return final_result
        
    except Exception as e:
        return f"เกิดข้อผิดพลาด: {str(e)}"

# สร้าง Gradio Interface
iface = gr.Interface(
    fn=process_audio,
    inputs=gr.Audio(type="filepath", label="อัพโหลดไฟล์เสียงของคุณ"),
    outputs=gr.Textbox(label="ผลการถอดความ"),
    title="ระบบถอดความเสียงภาษาไทยพร้อมแยกเสียงคนพูด",
    description="ใช้โมเดล Whisper สำหรับถอดความเสียงภาษาไทย และ pyannote.audio สำหรับแยกเสียงคนพูด"
)

# รัน app
iface.launch()