import gradio as gr
import whisperx
import os

# Загрузка модели один раз
model = whisperx.load_model("small", device="cpu")

# Основная функция транскрибации и диаризации
def transcribe_with_diarization(audio_path):
    result = model.transcribe(audio_path, diarize=True)
    segments = result["segments"]
    
    # Формируем вывод с возможностью редактирования
    text_output = ""
    for seg in segments:
        speaker = seg.get("speaker", "SPEAKER_00")
        start = seg["start"]
        end = seg["end"]
        text = seg["text"]
        text_output += f"[{speaker}] ({start:.2f} - {end:.2f}): {text}\n\n"
    
    return text_output

# Интерфейс
with gr.Blocks() as demo:
    gr.Markdown("# 📱🎤 Транскрибация и диаризация аудио")
    gr.Markdown("Загрузите аудиофайл, прослушайте сегменты, отредактируйте спикеров и выгрузите результат.")
    
    with gr.Row():
        audio_input = gr.Audio(source="upload", type="filepath", label="Загрузите аудиофайл")
    
    transcribe_btn = gr.Button("▶️ Транскрибировать")
    transcript_output = gr.Textbox(label="Результат транскрипции", lines=20)
    
    transcribe_btn.click(fn=transcribe_with_diarization, inputs=audio_input, outputs=transcript_output)

demo.launch()