Update app.py

app.py

@@ -3,90 +3,126 @@ 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") or os.getenv("hf_token") or ""
-    
+    hf_token = os.getenv("HF_TOKEN", "")
+
     with gr.Blocks() as app:
         gr.Markdown("<h1>Транскрипция и диаризация аудио</h1>")
-        gr.Markdown("Загрузите аудиофайл и нажмите **Транскрибировать**. После обработки вы сможете прослушать сегменты, отредактировать текст и присвоить имена спикерам.")
-        audio_input = gr.Audio(label="Аудиофайл", source="upload", type="filepath")
+        gr.Markdown(
+            "Загрузите аудиофайл (формат WAV/MP3), нажмите **Транскрибировать**, "
+            "отредактируйте результат и сохраните его."
+        )
+
+        # Убираем `source="upload"` — по умолчанию Audio позволяет загрузку
+        audio_input = gr.Audio(label="Аудиофайл", type="filepath")
         transcribe_btn = gr.Button("Транскрибировать")
-        
+
+        # Здесь будут динамически добавляться поля для редактирования
+        segment_container = gr.Column()
         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")
+        output_file = gr.File(label="Скачать .txt")
+
+        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 = 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 = 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
+            )
+
+            # 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)
+
+            # 3) Подготовка UI сегментов
             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))}
+
+            # Очищаем контейнер и добавляем новые поля
+            segment_container.clear()
+
+            # Поля для переименования спикеров
+            name_inputs = {}
+            with segment_container:
+                gr.Markdown("**Укажите имена спикеров:**")
+                for spk in speakers:
+                    name_inputs[spk] = gr.Textbox(
+                        label=f"Спикер {spk}",
+                        value=f"Спикер {spk}"
+                    )
+
+                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))
+
+            # Функция сохранения
+            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 idx, (speaker, _) in enumerate(transcripts):
-                        name = name_map.get(speaker, f"Спикер {speaker}")
-                        text = texts[idx]
-                        f.write(f"{name}: {text}\n")
+                    for spk, ti in text_inputs:
+                        text = kwargs[ti.label]
+                        f.write(f"{names[spk]}: {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)
+
+            # Создаем привязку кнопки сохранения
+            save_btn.click(
+                fn=save_result,
+                inputs=list(name_inputs.values()) + [ti for _, ti in text_inputs],
+                outputs=output_file
+            )
+
+        transcribe_btn.click(
+            fn=transcribe_with_diarization,
+            inputs=audio_input,
+            outputs=[]
+        )
+
+        app.launch(
+            server_name="0.0.0.0",
+            server_port=7860,
+            show_api=False
+        )
 
 if __name__ == "__main__":
     create_app()