changed speechlib to whisperx

app.py

@@ -3,43 +3,52 @@ import base64
 import os
 import json
 import streamlit as st
-from speechlib import Transcriptor
+import whisperx
+import torch
 
-def transcribe_audio(file, log_folder, language, modelSize, ACCESS_TOKEN, voices_folder, quantization):
-    transcriptor = Transcriptor(file, log_folder, language, modelSize, ACCESS_TOKEN, voices_folder, quantization)
-    return transcriptor.whisper()
-    
-def transform_transcript(transcript):
+def convert_segments_to_text(segments):
     result = []
-    for segment in transcript:
-        start_time, end_time, text, speaker = segment
-        result.append(f"{speaker} ({start_time:.1f} : {end_time:.1f}) : {text}")
+    for segment in segments:
+        speaker = segment['speaker']
+        start = segment['start']
+        end = segment['end']
+        text = segment['text']
+        formatted_text = f'{speaker} ({start} : {end}) : {text}'
+        result.append(formatted_text)
     return '\n'.join(result)
     
 st.title('Audio Transcription App')
+st.sidebar.title("Settings")
+# Sidebar inputs
+device = st.sidebar.selectbox("Device", ["cpu", "cuda"], index=1)
+batch_size = st.sidebar.number_input("Batch Size", min_value=1, value=16)
+compute_type = st.sidebar.selectbox("Compute Type", ["float16", "int8"], index=0)
 
 ACCESS_TOKEN = st.secrets["HF_TOKEN"]
 
 uploaded_file = st.file_uploader("Загрузите аудиофайл", type=["mp4", "wav", "m4a"])
 
 if uploaded_file is not None:
+    st.audio(uploaded_file)
     file_extension = uploaded_file.name.split(".")[-1]  # Получаем расширение файла
     temp_file_path = f"temp_file.{file_extension}"  # Создаем временное имя файла с правильным расширением
     
     with open(temp_file_path, "wb") as f:
         f.write(uploaded_file.getbuffer())
-        
-    log_folder = "logs"
-    language = "ru"
-    modelSize = os.getenv('WHISPER_MODEL_SIZE')
-    voices_folder = ""
-    quantization = False
     
     with st.spinner('Транскрибируем...'):
-        result = transcribe_audio(temp_file_path, log_folder, language, modelSize, ACCESS_TOKEN, voices_folder, quantization)
+        # Load model
+        model = whisperx.load_model("medium", device, compute_type=compute_type)
+        # Load and transcribe audio
+        audio = whisperx.load_audio(temp_file_path)
+        result = model.transcribe(audio, batch_size=batch_size, language="Russian")
+        # Load diarization model (replace YOUR_HF_TOKEN with actual token)
+        diarize_model = whisperx.DiarizationPipeline(use_auth_token=st.secrets["HF_TOKEN"], device=device)
+        diarize_segments = diarize_model(audio)
+        result = whisperx.assign_word_speakers(diarize_segments, result)
         
     st.write("Результат транскрибации:")
-    transcript = transform_transcript(result)
+    transcript = convert_segments_to_text(result)
     st.text(transcript)
     
     with st.spinner('Резюмируем...'):