Update app.py

app.py

@@ -1,128 +1,98 @@
 import os
+import tempfile
 import torch
 import whisperx
 from pyannote.audio import Pipeline
 import gradio as gr
-import torchaudio
 
-def create_app():
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    hf_token = os.getenv("HF_TOKEN", "")
-
-    with gr.Blocks() as app:
-        gr.Markdown("<h1>Транскрипция и диаризация аудио</h1>")
-        gr.Markdown(
-            "Загрузите аудиофайл (формат WAV/MP3), нажмите **Транскрибировать**, "
-            "отредактируйте результат и сохраните его."
-        )
-
-        # Убираем `source="upload"` — по умолчанию Audio позволяет загрузку
-        audio_input = gr.Audio(label="Аудиофайл", type="filepath")
-        transcribe_btn = gr.Button("Транскрибировать")
+# Загружаем модели один раз при старте
+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)
 
-        # Здесь будут динамически добавляться поля для редактирования
-        segment_container = gr.Column()
-        save_btn = gr.Button("Сохранить результат")
-        output_file = gr.File(label="Скачать .txt")
+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(
+        result["segments"],
+        asr_model.audio, 
+        asr_model.tokenizer, 
+        device
+    )
+    # 3) диаризация
+    diar = diar_model({"uri": "audio", "audio": audio_path})
+    # вплетаем спикер-теги
+    segments = whisperx.diarize(result["segments"], diar)
 
-        def transcribe_with_diarization(audio_path):
-            # 1) Транскрипция WhisperX с фиксированным языком "ru"
-            asr_model = whisperx.load_model("small", device, compute_type="float32")
-            audio_array = whisperx.load_audio(audio_path)
-            result = asr_model.transcribe(
-                audio_array,
-                batch_size=16,
-                language="ru"
-            )
-            align_model, metadata = whisperx.load_align_model(
-                language_code="ru", device=device
-            )
-            result = whisperx.align(
-                result["segments"],
-                align_model,
-                metadata,
-                audio_array,
-                device=device,
-                return_char_alignments=False
-            )
+    # 4) готовим текст для raw_output (объединяем, без спикеров)
+    full_text = "\n".join(f"[{seg['start']:.2f}-{seg['end']:.2f}] {seg['text']}"
+                          for seg in segments)
 
-            # 2) Диаризация Pyannote
-            diar_pipeline = Pipeline.from_pretrained(
-                "pyannote/speaker-diarization-3.1",
-                use_auth_token=hf_token
-            ).to(device)
-            diarization = diar_pipeline(audio_path)
-            result = whisperx.assign_word_speakers(diarization, result)
+    # сохраняем в 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()
 
-            # 3) Подготовка UI сегментов
-            segments = result["segments"]
-            speakers = sorted({seg["speaker"] for seg in segments})
+    return full_text, tmp.name
 
-            # Очищаем контейнер и добавляем новые поля
-            segment_container.clear()
+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(
+            "Загрузите аудиофайл, нажмите **Транскрибировать**, "
+            "прослушайте сегменты, отредактируйте текст и присвойте имена спикерам."
+        )
 
-            # Поля для переименования спикеров
-            name_inputs = {}
-            with segment_container:
-                gr.Markdown("**Укажите имена спикеров:**")
-                for spk in speakers:
-                    name_inputs[spk] = gr.Textbox(
-                        label=f"Спикер {spk}",
-                        value=f"Спикер {spk}"
-                    )
+        audio_input = gr.Audio(
+            label="Аудиофайл",
+            type="filepath"
+        )
 
-                gr.Markdown("---")
-                gr.Markdown("**Отредактируйте текст сегментов:**")
-                text_inputs = []
-                for i, seg in enumerate(segments):
-                    start, end = seg["start"], seg["end"]
-                    speaker = seg["speaker"]
-                    txt = seg["text"]
-                    # Срез аудио для сегмента
-                    seg_path = f"seg_{i}.wav"
-                    wave, sr = torchaudio.load(audio_path)
-                    torchaudio.save(
-                        seg_path,
-                        wave[:, int(start*sr):int(end*sr)],
-                        sr
-                    )
-                    with gr.Row():
-                        gr.Audio(value=seg_path, format="wav", label=None)
-                        ti = gr.Textbox(
-                            value=txt,
-                            label=f"{name_inputs[speaker].value}: {start:.1f}-{end:.1f}s",
-                            lines=2
-                        )
-                        text_inputs.append((speaker, ti))
+        transcribe_btn = gr.Button("▶️ Транскрибировать")
 
-            # Функция сохранения
-            def save_result(**kwargs):
-                # kwargs содержит сначала name_inputs, потом text_inputs
-                names = {spk: kwargs[f"Спикер {spk}"] for spk in speakers}
-                with open("result.txt", "w", encoding="utf-8") as f:
-                    for spk, ti in text_inputs:
-                        text = kwargs[ti.label]
-                        f.write(f"{names[spk]}: {text}\n")
-                return "result.txt"
+        # ПОЛЕ для «сырого» текста сразу после транскрипции
+        raw_output = gr.Textbox(
+            label="Результат транскрипции",
+            placeholder="Здесь появится текст после транскрибации",
+            lines=6
+        )
 
-            # Создаем привязку кнопки сохранения
-            save_btn.click(
-                fn=save_result,
-                inputs=list(name_inputs.values()) + [ti for _, ti in text_inputs],
-                outputs=output_file
-            )
+        save_btn = gr.Button("💾 Сохранить результат")
+        output_file = gr.File(
+            label="Скачать .txt",
+            file_count="single"
+        )
 
         transcribe_btn.click(
             fn=transcribe_with_diarization,
             inputs=audio_input,
-            outputs=[]
+            outputs=[raw_output, output_file]
         )
 
-        app.launch(
-            server_name="0.0.0.0",
-            server_port=7860,
-            show_api=False
-        )
+    return app
 
 if __name__ == "__main__":
-    create_app()
+    create_app().launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        inbrowser=False
+    )