Update app.py

app.py

@@ -1,67 +1,95 @@
 import os
+import tempfile
+import whisperx
 import torch
 import gradio as gr
-import whisperx
-from whisperx.diarize import DiarizationPipeline
 
 # 1. Устройство
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# 2. HuggingFace‑токен для диаризации (если нужен для закрытых пайплайнов)
+# 2. Загрузка моделей
+asr_model = whisperx.load_model("small", device)
 hf_token = os.getenv("HF_TOKEN", None)
+diarize_pipeline = whisperx.DiarizationPipeline(use_auth_token=hf_token, device=device)
 
-# 3. Загружаем WhisperX с принудительным float32
-model = whisperx.load_model(
-    "small",
-    device=device,
-    compute_type="float32"   # <- ключевая правка
-)
-
-# 4. Инициализируем пайплайн диаризации
-diarize_pipeline = DiarizationPipeline(use_auth_token=hf_token, device=device)
-
-def transcribe_with_diarization(audio_path):
-    # 5. Транскрипция, жестко на русском
-    result = model.transcribe(
-        audio_path,
-        language="ru"
-    )
-    # 6. Выравнивание таймингов
-    result = whisperx.align(
-        result["segments"],
-        model,
-        audio_path,
-        device=device
+def transcribe_and_prepare(audio_path):
+    # ASR (жёстко русский)
+    result = asr_model.transcribe(audio_path, language="ru")
+
+    # Alignment
+    aligned = whisperx.align(
+        result["segments"], asr_model, audio_path, device=device
     )
-    # 7. Диаризация
+
+    # Diarization
     diarization = diarize_pipeline(audio_path)
-    # 8. Объединение сегментов с метками спикеров
-    segments = whisperx.diarize(
-        result["segments"],
-        diarization
+    segments = whisperx.diarize(aligned, diarization)
+
+    # Подготовка для UI: возвращаем список dict-ов
+    ui_data = []
+    for i, seg in enumerate(segments):
+        ui_data.append({
+            "index": i,
+            "speaker": seg["speaker"],
+            "start": f"{seg['start']:.2f}",
+            "end": f"{seg['end']:.2f}",
+            "text": seg["text"]
+        })
+    return ui_data
+
+def generate_download(ui_data):
+    # Формируем итоговый TXT
+    lines = []
+    for row in ui_data:
+        lines.append(f"[{row['speaker']}] ({row['start']}-{row['end']}): {row['text']}")
+    txt = "\n".join(lines)
+    path = os.path.join(tempfile.gettempdir(), "transcript.txt")
+    with open(path, "w", encoding="utf-8") as f:
+        f.write(txt)
+    return path
+
+# 3. Интерфейс
+with gr.Blocks(css="""
+.gradio-container { max-width: 900px; margin: auto; }
+@media (max-width: 600px) {
+  .gradio-container { padding: 0 10px; }
+}
+""") as demo:
+
+    gr.Markdown("## 🎤 Транскрибация и диаризация аудио (русский)")
+    audio_in = gr.Audio(label="Загрузите аудио", type="filepath")
+    btn = gr.Button("Запустить транскрибацию")
+
+    # Таблица сегментов для ручной правки
+    table = gr.Dataframe(
+        headers=["index","speaker","start","end","text"],
+        datatype=["number","text","text","text","text"],
+        interactive=True,
+        row_count=(1, None),
+        col_count=5,
+        wrap=True,
+        label="Сегменты (можно править спикера и текст)"
     )
-    # 9. Формирование выходного текста
-    out = ""
-    for seg in segments:
-        spk = seg["speaker"]
-        txt = seg["text"].strip()
-        out += f"[{spk}]: {txt}\n"
-    return out
-
-# 10. Web‑интерфейс Gradio
-app = gr.Blocks(css="""
-  /* мобильная адаптация */
-  @media(max-width:600px) {
-    .gradio-container { padding: 0 1rem; }
-  }
-""")
-
-with app:
-    gr.Markdown("## Транскрибация и диаризация (русский)\nЗагрузите аудио, нажмите «Транскрибировать»")
-    audio_input = gr.Audio(label="Аудиофайл", type="filepath")
-    btn = gr.Button("Транскрибировать")
-    output = gr.Textbox(label="Результат", lines=20)
-    btn.click(transcribe_with_diarization, inputs=audio_input, outputs=output)
+
+    download_btn = gr.Button("Скачать итоговый TXT")
+    download_txt = gr.File(label="Итоговый файл")
+
+    # Связываем
+    btn.click(fn=transcribe_and_prepare, inputs=[audio_in], outputs=[table])
+    download_btn.click(fn=generate_download, inputs=[table], outputs=[download_txt])
+
+    # Плейер для выбранного сегмента
+    with gr.Row():
+        idx_in = gr.Number(value=0, label="Номер сегмента для прослушивания")
+        play_btn = gr.Button("▶️ Прослушать сегмент")
+        player = gr.Audio(label="Плеер сегмента")
+
+    def play_segment(audio_path, ui_data, idx):
+        seg = ui_data[int(idx)]
+        start, end = float(seg["start"]), float(seg["end"])
+        return {"filepath": audio_path, "start_time": start, "end_time": end}
+
+    play_btn.click(fn=play_segment, inputs=[audio_in, table, idx_in], outputs=[player])
 
 if __name__ == "__main__":
-    app.launch()
+    demo.launch()