Update app.py

app.py

@@ -1,95 +1,92 @@
 import os
-import tempfile
-import whisperx
 import torch
+import whisperx
+from pyannote.audio import Pipeline
 import gradio as gr
 
-# 1. Устройство
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# 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)
-
-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
-    )
-
-    # Diarization
-    diarization = diarize_pipeline(audio_path)
-    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="Сегменты (можно править спикера и текст)"
-    )
-
-    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])
+def create_app():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    hf_token = os.getenv("HF_TOKEN") or os.getenv("hf_token") or ""
+    
+    with gr.Blocks() as app:
+        gr.Markdown("<h1>Транскрипция и диаризация аудио</h1>")
+        gr.Markdown("Загрузите аудиофайл и нажмите **Транскрибировать**. После обработки вы сможете прослушать сегменты, отредактировать текст и присвоить имена спикерам.")
+        audio_input = gr.Audio(label="Аудиофайл", source="upload", type="filepath")
+        transcribe_btn = gr.Button("Транскрибировать")
+        
+        save_btn = gr.Button("Сохранить результат")
+        output_file = gr.File(label="Скачайте результат (.txt)")
+        
+        @gr.render(inputs=[audio_input], triggers=[transcribe_btn])
+        def process(audio_path):
+            if not audio_path:
+                return
+            # 1. WhisperX transcription
+            model = whisperx.load_model("small", device, compute_type="float32")
+            audio_array = whisperx.load_audio(audio_path)
+            result = model.transcribe(audio_array, batch_size=16, language="ru")
+            model_a, metadata = whisperx.load_align_model(language_code="ru", device=device)
+            result = whisperx.align(result["segments"], model_a, metadata, audio_array, device=device, return_char_alignments=False)
+            # 2. Speaker diarization
+            pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1", use_auth_token=hf_token)
+            pipeline.to(device)
+            diarization = pipeline(audio_path)
+            result = whisperx.assign_word_speakers(diarization, result)
+            segments = result["segments"]
+            # Unique speakers
+            speakers = sorted({seg["speaker"] for seg in segments})
+            
+            # Input fields for speaker names
+            if speakers:
+                gr.Markdown("**Имена спикеров:**")
+            speaker_name_inputs = []
+            for spk in speakers:
+                tb = gr.Textbox(label=f"Спикер {spk}", value=f"Спикер {spk}", interactive=True, key=f"name_{spk}")
+                speaker_name_inputs.append(tb)
+            
+            # Load audio for slicing segments
+            try:
+                import torchaudio
+                waveform, sample_rate = torchaudio.load(audio_path)
+            except Exception:
+                waveform, sample_rate = None, None
+            
+            transcripts = []
+            transcript_text_inputs = []
+            # Render each segment
+            for i, seg in enumerate(segments):
+                speaker = seg["speaker"]
+                text = seg["text"]
+                transcripts.append((speaker, text))
+                start, end = seg["start"], seg["end"]
+                seg_audio_path = audio_path
+                if waveform is not None and sample_rate is not None:
+                    start_idx = int(start * sample_rate)
+                    end_idx = int(end * sample_rate)
+                    segment_waveform = waveform[:, start_idx:end_idx]
+                    seg_audio_path = f"segment_{i}.wav"
+                    torchaudio.save(seg_audio_path, segment_waveform, sample_rate)
+                with gr.Row():
+                    gr.Audio(value=seg_audio_path, format="audio/wav", show_label=False, interactive=False, key=f"audio_{i}")
+                    tb_seg = gr.Textbox(value=text, lines=2, label=f"Спикер {speaker}", key=f"text_{i}", interactive=True)
+                    transcript_text_inputs.append(tb_seg)
+            
+            # Define save function
+            def save_func(*args):
+                names = list(args[:len(speakers)])
+                texts = list(args[len(speakers):])
+                name_map = {speakers[j]: names[j] for j in range(len(speakers))}
+                with open("result.txt", "w", encoding="utf-8") as f:
+                    for idx, (speaker, _) in enumerate(transcripts):
+                        name = name_map.get(speaker, f"Спикер {speaker}")
+                        text = texts[idx]
+                        f.write(f"{name}: {text}\n")
+                return "result.txt"
+            
+            if speakers or transcripts:
+                save_btn.click(save_func, inputs=speaker_name_inputs + transcript_text_inputs, outputs=output_file)
+        
+        app.launch(server_name="0.0.0.0", server_port=7860)
 
 if __name__ == "__main__":
-    demo.launch()
+    create_app()