Update app.py

app.py

@@ -5,25 +5,28 @@ import json
 import os
 import requests
 import torch
-import numpy
+import numpy as np
+import librosa
 from datetime import timedelta
 from pyannote.audio import Pipeline
-from huggingface_hub import login
-
-# --- Fix for OMP Error ---
-os.environ["OMP_NUM_THREADS"] = "1"
+from huggingface_hub import login, hf_hub_download
 
 # --- Safe Globals for PyTorch 2.6+ ---
 try:
     from pyannote.audio.core.task import Specifications, Problem, Resolution
+    from pyannote.audio.core.model import Model
+    from pyannote.audio.pipelines.speaker_diarization import SpeakerDiarization
     
     torch.serialization.add_safe_globals([
         torch.torch_version.TorchVersion,
         Specifications,
         Problem,
         Resolution,
-        numpy.dtype,
+        Model,
+        SpeakerDiarization,
+        np.dtype,
         torch.nn.modules.container.ModuleList,
+        np.core.multiarray.scalar
     ])
 except Exception as e:
     print(f"Safe Globals Warning: {e}")
@@ -125,96 +128,104 @@ if uploaded_file:
             if not ACTIVE_HF_TOKEN or "PASTE_YOUR_HF_TOKEN" in ACTIVE_HF_TOKEN:
                 st.error("Please provide a valid Hugging Face Token.")
             else:
-                status_container = st.empty()
-                with status_container.container():
-                    st.write("🔄 **Processing Started...**")
-                    
-                    # Authenticate FIRST - this is critical
+                with st.spinner("Processing... This may take a moment."):
+                    # Authenticate Globally
                     login(token=ACTIVE_HF_TOKEN)
                     
                     with open("temp_input", "wb") as f:
                         f.write(uploaded_file.getbuffer())
                     
-                    st.write("🎵 **Extracting Audio (WAV)...**")
+                    # Convert to strict WAV using FFmpeg
                     subprocess.run([
                         "ffmpeg", "-i", "temp_input", 
                         "-vn", "-acodec", "pcm_s16le", "-ar", "16000", "-ac", "1", 
                         "temp_audio.wav", "-y"
-                    ], capture_output=True)
+                    ])
                     
-                    # 1. Diarization - CRITICAL FIX
+                    # 1. Diarization with Librosa (Robust Audio Loading)
                     st.write("🗣️ **Running Speaker Diarization...**")
                     diarization = None
-                    speaker_turns = []
-                    
                     try:
-                        # Load pipeline WITHOUT any token parameter (relies on login())
-                        pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1")
+                        # Load Config Manually
+                        config_path = hf_hub_download(
+                            repo_id="pyannote/speaker-diarization-3.1", 
+                            filename="config.yaml", 
+                            token=ACTIVE_HF_TOKEN
+                        )
+                        pipeline = Pipeline.from_pretrained(config_path)
                         
+                        # Load Audio with Librosa (More robust than torchaudio in some containers)
+                        y, sr = librosa.load("temp_audio.wav", sr=16000)
+                        waveform = torch.tensor(y).unsqueeze(0) # Add channel dim
+                        
+                        # Move to GPU if available
                         if torch.cuda.is_available():
                             st.write("🚀 Using GPU for Diarization")
                             pipeline.to(torch.device("cuda"))
-                        
-                        # Run diarization
-                        st.write("⏳ Processing audio (this takes a moment)...")
-                        diarization = pipeline("temp_audio.wav")
-                        
-                        # Debug: Check what we got back
-                        st.write(f"DEBUG: Diarization type = {type(diarization)}")
-                        st.write(f"DEBUG: Has itertracks? {hasattr(diarization, 'itertracks')}")
-                        
-                        # Extract speaker turns
-                        if hasattr(diarization, 'itertracks'):
-                            for turn, _, speaker_id in diarization.itertracks(yield_label=True):
-                                speaker_turns.append({
-                                    "start": turn.start, 
-                                    "end": turn.end, 
-                                    "speaker": speaker_id
-                                })
-                            st.write(f"✅ Found {len(speaker_turns)} speaker segments")
+                            waveform = waveform.to(torch.device("cuda"))
                             
-                            # Show first few turns for debugging
-                            if speaker_turns:
-                                st.write("Sample speaker turns:")
-                                for turn in speaker_turns[:3]:
-                                    st.write(f"  {turn['speaker']}: {turn['start']:.1f}s - {turn['end']:.1f}s")
+                        # Run Pipeline on Tensor directly
+                        diarization_output = pipeline({"waveform": waveform, "sample_rate": 16000})
+                        
+                        # Handle Output Wrapper
+                        if isinstance(diarization_output, tuple):
+                            diarization = diarization_output[0]
                         else:
-                            st.warning("⚠️ Diarization object has no itertracks method")
+                            diarization = diarization_output
                             
+                        # Extract Annotation
+                        if not hasattr(diarization, "itertracks"):
+                             if hasattr(diarization_output, "annotation"):
+                                 diarization = diarization_output.annotation
+                             elif hasattr(diarization_output, "get"):
+                                 diarization = diarization_output.get("annotation", diarization_output)
+
                     except Exception as e:
-                        st.error(f"Diarization Error: {str(e)}")
-                        import traceback
-                        st.code(traceback.format_exc())
+                        st.error(f"Diarization Error: {e}")
+                        diarization = None
                     
                     # 2. Whisper Transcription
                     st.write("📝 **Transcribing with Whisper...**")
                     device = "cuda" if torch.cuda.is_available() else "cpu"
-                    model = whisper.load_model("base", device=device)
+                    model = whisper.load_model("medium", device=device) 
                     result = model.transcribe("temp_audio.wav", word_timestamps=True)
                     
                     # 3. Alignment
                     st.write("🔗 **Aligning Speakers...**")
                     final_segments = []
+                    speaker_turns = []
                     
+                    if diarization:
+                        try:
+                            iterator = None
+                            if hasattr(diarization, 'itertracks'):
+                                iterator = diarization.itertracks(yield_label=True)
+                            
+                            if iterator:
+                                for turn, _, speaker_id in iterator:
+                                    speaker_turns.append({"start": turn.start, "end": turn.end, "speaker": speaker_id})
+                                st.write(f"✅ Found {len(speaker_turns)} speaker turns.")
+                            else:
+                                st.warning("⚠️ Pipeline ran but returned no iterable tracks.")
+                        except Exception as e:
+                            st.error(f"Error iterating tracks: {e}")
+
                     for segment in result['segments']:
                         mid_time = (segment['start'] + segment['end']) / 2
                         speaker = "Unknown"
                         
                         if speaker_turns:
-                            # Match speaker by midpoint
+                            # Match speaker
                             for turn in speaker_turns:
                                 if turn["start"] <= mid_time <= turn["end"]:
                                     speaker = turn["speaker"]
                                     break
                             
-                            # Fallback: find closest speaker within 1 second
+                            # Fallback distance matching
                             if speaker == "Unknown":
-                                best_dist = 1.0
+                                best_dist = 1.0 
                                 for turn in speaker_turns:
-                                    dist = min(
-                                        abs(turn["start"] - mid_time), 
-                                        abs(turn["end"] - mid_time)
-                                    )
+                                    dist = min(abs(turn["start"] - mid_time), abs(turn["end"] - mid_time))
                                     if dist < best_dist:
                                         best_dist = dist
                                         speaker = turn["speaker"]
@@ -228,16 +239,15 @@ if uploaded_file:
                         })
                     
                     st.session_state.transcript = final_segments
-                    st.success(f"✓ Complete! Transcribed {len(final_segments)} segments")
+                    st.success("Complete!")
 
     if "transcript" in st.session_state:
         st.divider()
-        
         with st.expander("Transcript Preview (Diarized)"):
             for seg in st.session_state.transcript[:20]: 
                 st.markdown(f"**{seg['speaker']}:** {seg['text']}")
         
-        brief = st.text_area("Creative Brief", placeholder="e.g. Focus on the yeast story, remove the interviewer.")
+        brief = st.text_area("Creative Brief", placeholder="e.g. Focus on the yeast story.")
         
         if st.button("Step 2: Create EDL"):
             if not ACTIVE_GEMINI_KEY: