Update app.py

app.py

@@ -1,41 +1,98 @@
 import streamlit as st
 from pyannote.audio import Pipeline
-from transformers import pipeline
 import whisper
+import tempfile
+import os
+import torch
+from transformers import pipeline as tf_pipeline
 
-# Title
-st.title("Multi-Speaker Audio Analyzer")
+# Cache the model loading using streamlit
+@st.cache_resource
+def load_models():
+    try:
+        # Load diarization model efficiently
+        diarization = Pipeline.from_pretrained(
+            "pyannote/speaker-diarization",
+            use_auth_token=st.secrets["hf_token"]
+        )
+        
+        # Load smaller whisper model for faster processing
+        transcriber = whisper.load_model("base")
+        
+        # Load efficient summarizer
+        summarizer = tf_pipeline(
+            "summarization", 
+            model="facebook/bart-large-cnn",
+            device=0 if torch.cuda.is_available() else -1
+        )
+        
+        return diarization, transcriber, summarizer
+    except Exception as e:
+        st.error(f"Error loading models: {str(e)}")
+        return None, None, None
 
-# Upload Audio File
-uploaded_file = st.file_uploader("Upload an audio file (MP3/WAV)", type=["mp3", "wav"])
+def process_audio(audio_file, max_duration=300):  # limit to 5 minutes initially
+    try:
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
+            tmp.write(audio_file.getvalue())
+            tmp_path = tmp.name
 
-# Process Button
-if uploaded_file:
-    st.audio(uploaded_file, format='audio/wav')
+        # Get cached models
+        diarization, transcriber, summarizer = load_models()
+        if not all([diarization, transcriber, summarizer]):
+            return "Model loading failed"
 
-    # Load pre-trained models
-    diarization_pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization")
-    transcription_model = whisper.load_model("base")
-    summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
+        # Process with progress bar
+        with st.spinner("Identifying speakers..."):
+            diarization_result = diarization(tmp_path)
+        
+        with st.spinner("Transcribing audio..."):
+            transcription = transcriber.transcribe(tmp_path)
+            
+        with st.spinner("Generating summary..."):
+            summary = summarizer(transcription["text"], max_length=130, min_length=30)
 
-    # Perform Speaker Diarization
-    st.write("Processing Speaker Diarization...")
-    diarized_output = diarization_pipeline(uploaded_file)
+        # Cleanup
+        os.unlink(tmp_path)
+        
+        return {
+            "diarization": diarization_result,
+            "transcription": transcription["text"],
+            "summary": summary[0]["summary_text"]
+        }
+        
+    except Exception as e:
+        st.error(f"Error processing audio: {str(e)}")
+        return None
 
-    # Perform Speech-to-Text Transcription
-    st.write("Transcribing Audio...")
-    transcription = transcription_model.transcribe(uploaded_file)
+def main():
+    st.title("Multi-Speaker Audio Analyzer")
+    st.write("Upload an audio file (MP3/WAV) up to 5 minutes long for best performance")
 
-    # Generate Summary
-    st.write("Generating Summary...")
-    summary = summarizer(transcription["text"])
+    uploaded_file = st.file_uploader("Choose a file", type=["mp3", "wav"])
 
-    # Display Outputs
-    st.write("Speaker-Diarized Transcript:")
-    st.text(diarized_output)
+    if uploaded_file:
+        st.audio(uploaded_file, format='audio/wav')
+        
+        if st.button("Analyze Audio"):
+            results = process_audio(uploaded_file)
+            
+            if results:
+                # Display results in tabs
+                tab1, tab2, tab3 = st.tabs(["Speakers", "Transcription", "Summary"])
+                
+                with tab1:
+                    st.write("Speaker Segments:")
+                    for turn, _, speaker in results["diarization"].itertracks(yield_label=True):
+                        st.write(f"{speaker}: {turn.start:.1f}s → {turn.end:.1f}s")
+                
+                with tab2:
+                    st.write("Transcription:")
+                    st.write(results["transcription"])
+                
+                with tab3:
+                    st.write("Summary:")
+                    st.write(results["summary"])
 
-    st.write("Full Transcription:")
-    st.text(transcription["text"])
-
-    st.write("Summary:")
-    st.text(summary[0]['summary_text'])
+if __name__ == "__main__":
+    main()