Update app.py

app.py

@@ -1,9 +1,25 @@
 import os
 import numpy as np
+import torch
+
+# --- PYTORCH 2.6+ SECURITY & COMPATIBILITY PATCHES ---
+# 1. Allow WhisperX/Pyannote globals for model loading
+try:
+    from omegaconf.listconfig import ListConfig
+    from omegaconf.dictconfig import DictConfig
+    # Expanded safe globals to include classes often used in diarization checkpoints
+    torch.serialization.add_safe_globals([
+        ListConfig, 
+        DictConfig,
+        torch.nn.modules.container.ModuleList,
+        np.dtype,
+        np._core.multiarray.scalar if hasattr(np, '_core') else np.core.multiarray.scalar
+    ])
+except Exception as e:
+    # Use a generic print or pass to avoid startup crashes if classes are missing
+    print(f"Safe Globals Warning: {e}")
 
-# --- NUMPY 2.0 PATCH ---
-# This must run before any other library imports to prevent crashes
-# with pyannote/whisperx which might expect the old 'np.NaN' attribute.
+# 2. Fix NumPy 2.0+ attribute removal (required for older pyannote internals)
 if not hasattr(np, 'NaN'):
     np.NaN = np.nan
 
@@ -11,14 +27,12 @@ import streamlit as st
 import subprocess
 import json
 import requests
-import torch
 import whisperx
 import gc
 import pandas as pd
 from datetime import timedelta
 
 # --- Configuration & Tokens ---
-# Priority: Secret > Hardcoded
 HARDCODED_HF_TOKEN = "PASTE_YOUR_HF_TOKEN_HERE"
 HARDCODED_GEMINI_KEY = ""
 
@@ -100,8 +114,7 @@ with st.sidebar:
     fps = st.number_input("Timeline FPS", value=25)
     
     st.header("Model Settings")
-    # T4 has 16GB VRAM, large-v2 works well
-    model_size = st.selectbox("Whisper Model", ["large-v2", "medium"], index=0)
+    model_size = st.selectbox("Whisper Model", ["large-v2", "medium", "base"], index=0)
     num_speakers = st.number_input("Speakers (0=Auto)", min_value=0, value=0)
     
     st.divider()
@@ -114,7 +127,6 @@ with st.sidebar:
 uploaded_file = st.file_uploader("Upload Video/Audio Clip", type=["mp4", "m4a", "wav", "mp3", "mov"])
 
 if uploaded_file:
-    # --- Step 1: Technical Processing ---
     if "transcript" not in st.session_state:
         if st.button("Step 1: Transcribe & Diarize"):
             if not ACTIVE_HF_TOKEN or "PASTE_YOUR_HF_TOKEN" in ACTIVE_HF_TOKEN:
@@ -124,12 +136,10 @@ if uploaded_file:
                 with status_container.container():
                     st.write("🔄 **Processing Started...**")
                     
-                    # Save local temp file
                     with open("temp_input", "wb") as f:
                         f.write(uploaded_file.getbuffer())
                     
-                    st.write("🎵 **Extracting Audio (WAV)...**")
-                    # WhisperX prefers 16k mono wav
+                    st.write("🎵 **Extracting Audio...**")
                     subprocess.run([
                         "ffmpeg", "-i", "temp_input", 
                         "-vn", "-acodec", "pcm_s16le", "-ar", "16000", "-ac", "1", 
@@ -138,38 +148,33 @@ if uploaded_file:
                     
                     try:
                         device = "cuda" if torch.cuda.is_available() else "cpu"
-                        if device == "cpu":
-                            st.warning("⚠️ No GPU detected. WhisperX will be very slow.")
-                            
-                        st.write(f"🚀 **Loading WhisperX on {device}...**")
+                        st.write(f"🚀 **Running WhisperX on {device}...**")
                         
                         # 1. Transcribe
-                        # Use float16 for GPU, int8 for CPU
                         compute_type = "float16" if device == "cuda" else "int8"
-                        
                         model = whisperx.load_model(model_size, device, compute_type=compute_type)
                         
                         st.write("📝 **Transcribing...**")
                         audio = whisperx.load_audio("temp_audio.wav")
                         result = model.transcribe(audio, batch_size=16)
                         
-                        # Cleanup VRAM
+                        # Memory management
+                        del model
                         gc.collect()
                         torch.cuda.empty_cache()
-                        del model
                         
-                        # 2. Align (Improves timestamp accuracy for diarization)
+                        # 2. Align
                         st.write("⏱️ **Aligning Audio...**")
                         model_a, metadata = whisperx.load_align_model(language_code=result["language"], device=device)
                         result = whisperx.align(result["segments"], model_a, metadata, audio, device, return_char_alignments=False)
                         
-                        # Cleanup VRAM
+                        del model_a
                         gc.collect()
                         torch.cuda.empty_cache()
-                        del model_a
                         
                         # 3. Diarize
                         st.write("🗣️ **Diarizing Speakers...**")
+                        # Pass token for gated diarization models
                         diarize_model = whisperx.DiarizationPipeline(use_auth_token=ACTIVE_HF_TOKEN, device=device)
                         
                         diarize_kwargs = {}
@@ -178,13 +183,11 @@ if uploaded_file:
                             
                         diarize_segments = diarize_model(audio, **diarize_kwargs)
                         
-                        # 4. Assign Speakers to Words
+                        # 4. Final Merge
                         st.write("🔗 **Merging Transcripts...**")
                         final_result = whisperx.assign_word_speakers(diarize_segments, result)
                         
-                        # Format for Gemini
                         processed_segments = []
-                        # WhisperX structure is slightly different, it returns 'segments' list
                         for segment in final_result["segments"]:
                             processed_segments.append({
                                 "speaker": segment.get("speaker", "Unknown"),