Update app.py

app.py

@@ -1,98 +1,100 @@
 import os
 import tempfile
+import datetime
+
+import gradio as gr
 import torch
 import whisperx
-from pyannote.audio import Pipeline
-import gradio as gr
+from whisperx.diarize import DiarizationPipeline
 
-# Загружаем модели один раз при старте
+# Определяем устройство: CUDA если доступна, иначе CPU
 device = "cuda" if torch.cuda.is_available() else "cpu"
-asr_model = whisperx.load_model("small", device)  # модель WhisperX
-diar_model = Pipeline.from_pretrained(
-    "pyannote/speaker-diarization-3.1",
-    use_auth_token=os.getenv("HF_TOKEN", "")
-).to(device)
+
+# Загружаем модель WhisperX с compute_type="float32" и русским языком
+asr_model = whisperx.load_model(
+    "small",
+    device=device,
+    compute_type="float32"  # принудительная настройка, убирает float16
+)
+
+# Загружаем модель выравнивания для русского
+align_model, metadata = whisperx.load_align_model(
+    language_code="ru",
+    device=device
+)
+
+# Инициализируем пайплайн диаризации (Pyannote) с токеном HF
+hf_token = os.getenv("HUGGINGFACEHUB_API_TOKEN", None)
+diarization_pipeline = DiarizationPipeline(
+    use_auth_token=hf_token,
+    device=device
+)
 
 def transcribe_with_diarization(audio_path):
-    # 1) транскрипция
-    result = asr_model.transcribe(
-        audio_path,
-        language="ru",            # фиксируем русский, чтобы не тратить время на детект
-        compute_type="float32",   # CPU-friendly
-        diarize=False
-    )
-    # 2) выравнивание (alignment)
-    result = whisperx.align(
+    # 1) ASR без детекции языка (принудительно ru)
+    result = asr_model.transcribe(audio_path, language="ru")
+
+    # 2) Выравнивание субтитров по аудио
+    aligned = whisperx.align(
         result["segments"],
-        asr_model.audio, 
-        asr_model.tokenizer, 
+        align_model,
+        metadata,
+        audio_path,
         device
     )
-    # 3) диаризация
-    diar = diar_model({"uri": "audio", "audio": audio_path})
-    # вплетаем спикер-теги
-    segments = whisperx.diarize(result["segments"], diar)
-
-    # 4) готовим текст для raw_output (объединяем, без спикеров)
-    full_text = "\n".join(f"[{seg['start']:.2f}-{seg['end']:.2f}] {seg['text']}"
-                          for seg in segments)
-
-    # сохраняем в temp-файл для кнопки "Сохранить"
-    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".txt", mode="w", encoding="utf-8")
-    for seg in segments:
-        tmp.write(f"{seg['speaker']}: {seg['start']:.2f}-{seg['end']:.2f}\t{seg['text']}\n")
-    tmp.close()
-
-    return full_text, tmp.name
-
-def create_app():
-    with gr.Blocks(
-        title="Транскрипция и диаризация аудио",
-        css="""
-        @media(max-width:600px) {
-            .gradio-container { padding: 0.5rem; }
-            .gr-button { width: 100% !important; }
-        }
-        """
-    ) as app:
-        gr.Markdown("# 🎙️ Транскрипция и диаризация аудио")
-        gr.Markdown(
-            "Загрузите аудиофайл, нажмите **Транскрибировать**, "
-            "прослушайте сегменты, отредактируйте текст и присвойте имена спикерам."
-        )
-
-        audio_input = gr.Audio(
-            label="Аудиофайл",
-            type="filepath"
-        )
-
-        transcribe_btn = gr.Button("▶️ Транскрибировать")
-
-        # ПОЛЕ для «сырого» текста сразу после транскрипции
-        raw_output = gr.Textbox(
-            label="Результат транскрипции",
-            placeholder="Здесь появится текст после транскрибации",
-            lines=6
-        )
-
-        save_btn = gr.Button("💾 Сохранить результат")
-        output_file = gr.File(
-            label="Скачать .txt",
-            file_count="single"
-        )
-
-        transcribe_btn.click(
-            fn=transcribe_with_diarization,
-            inputs=audio_input,
-            outputs=[raw_output, output_file]
-        )
-
-    return app
 
-if __name__ == "__main__":
-    create_app().launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,
-        inbrowser=False
+    # 3) Диаризация
+    diarization = diarization_pipeline(audio_path)
+
+    # 4) Объединяем текстовые сегменты и спикеров
+    merged = whisperx.merge_text_with_diarization(
+        aligned["segments"],
+        diarization["segments"]
     )
+
+    # 5) Формируем текст для вывода
+    lines = []
+    for seg in merged:
+        spk = seg.get("speaker", "Speaker")
+        txt = seg.get("text", "").strip()
+        lines.append(f"[{spk}] {txt}")
+    return "\n".join(lines)
+
+def export_to_txt(text):
+    # Сохраняем результат во временный файл и возвращаем путь
+    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+    filename = f"transcript_{timestamp}.txt"
+    path = os.path.join(tempfile.gettempdir(), filename)
+    with open(path, "w", encoding="utf-8") as f:
+        f.write(text)
+    return path
+
+# Собираем интерфейс Gradio
+app = gr.Blocks(title="🎙️ DiarAI: Транскрибация и диаризация (RU)")
+
+with app:
+    gr.Markdown("""
+    ## Транскрибация и диаризация (русский язык)
+    - Фиксированный язык распознавания: **ru** для повышения скорости.
+    - Диаризация спикеров через Pyannote.
+    """)
+
+    audio_input = gr.Audio(type="filepath", label="Загрузите аудио (только RU)")
+    transcribe_btn = gr.Button("▶️ Транскрибировать")
+    output_txt = gr.Textbox(label="Результат транскрипции", lines=20)
+    save_btn = gr.Button("💾 Экспорт в .txt")
+    download_file = gr.File(label="Скачать результат")
+
+    transcribe_btn.click(
+        fn=transcribe_with_diarization,
+        inputs=audio_input,
+        outputs=output_txt
+    )
+    save_btn.click(
+        fn=export_to_txt,
+        inputs=output_txt,
+        outputs=download_file
+    )
+
+if __name__ == "__main__":
+    app.launch()