Update app.py

app.py

@@ -112,6 +112,13 @@ def load_pipeline():
         st.code(traceback.format_exc())
         raise e
 
+@st.cache_resource
+def load_speechbrain_model():
+    st.info("Loading SpeechBrain model...")
+    classifier = EncoderClassifier.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb")
+    st.success("SpeechBrain model loaded successfully")
+    return classifier
+
 # Sidebar
 with st.sidebar:
     st.header("Settings")
@@ -124,6 +131,7 @@ with st.sidebar:
 # Main content
 tab1, tab2, tab3 = st.tabs(["Upload & Process", "Results", "Visualization"])
 
+
 with tab1:
     uploaded_file = st.file_uploader("Choose an audio file", type=['wav', 'mp3', 'flac'])
     
@@ -138,16 +146,17 @@ with tab1:
                 check_hf_api()
                 verify_model_files()
                 pipeline = load_pipeline()
+                speechbrain_model = load_speechbrain_model()
             else:
                 st.stop()
 
             with st.status("Processing audio...", expanded=True) as status:
                 progress_bar = st.progress(0)
                 
-                def progress_hook(step: int, total: int, stage: str):
-                    if total > 0:
-                        progress_percentage = min(step / total, 1.0)
-                        status.update(label=f"Processing: {stage} ({progress_percentage:.1%})", state="running")
+                def progress_hook(**kwargs):
+                    if 'completed' in kwargs and 'total' in kwargs:
+                        progress_percentage = min(kwargs['completed'] / kwargs['total'], 1.0)
+                        status.update(label=f"Processing: {progress_percentage:.1%} complete", state="running")
                         progress_bar.progress(progress_percentage)
 
                 # Run the pipeline on the audio file
@@ -171,6 +180,11 @@ with tab1:
                 rttm_line = f"SPEAKER {os.path.basename(tmp_path)} 1 {turn.start:.3f} {turn.duration:.3f} <NA> <NA> {speaker} <NA> <NA>\n"
                 rttm_content += rttm_line
 
+            # Use SpeechBrain for speaker embedding (optional)
+            waveform, sample_rate = torchaudio.load(tmp_path)
+            embeddings = speechbrain_model.encode_batch(waveform)
+            st.success("Speaker embeddings generated successfully")
+
         except Exception as e:
             st.error(f"An error occurred: {str(e)}")
             st.error("Error details:")