Update app.py

app.py

@@ -2,50 +2,66 @@ import os
 import numpy as np
 import torch
 import typing
+import torchaudio
 
-# --- CRITICAL: AGGRESSIVE PYTORCH SECURITY OVERRIDE ---
-# The error "Weights only load failed" means weights_only=True is active.
-# We must intercept torch.load and FORCE it to False for pyannote/whisperx to work.
-_original_torch_load = torch.load
+# --- CRITICAL ENVIRONMENT FIXES ---
+# 1. Fix for Hugging Face millicore OMP_NUM_THREADS error
+if os.environ.get("OMP_NUM_THREADS", "").endswith("m"):
+    os.environ["OMP_NUM_THREADS"] = "1"
 
+# 2. Force Torchaudio Backend
+try:
+    if "ffmpeg" in torchaudio.list_audio_backends():
+        torchaudio.set_audio_backend("ffmpeg")
+except Exception:
+    pass
+
+# 3. GLOBAL PYTORCH 2.6+ SECURITY BYPASS (MONKEYPATCH)
+_original_torch_load = torch.load
 def patched_torch_load(*args, **kwargs):
-    # Debug print to verify interception (Check Logs tab in HF)
-    print(f"DEBUG: Intercepted torch.load call. Target: {args[0] if args else 'Unknown'}")
+    # Debug print to see what's being loaded
+    # print(f"DEBUG: Intercepted torch.load call. Target: {args[0] if args else 'Unknown'}")
     
     # FORCE Disable security check
     kwargs['weights_only'] = False 
-    
     return _original_torch_load(*args, **kwargs)
 
-# Apply the patch
 torch.load = patched_torch_load
 print("DEBUG: torch.load has been monkeypatched to allow all globals.")
 
-# --- Secondary Backup: Safe Globals Whitelist ---
+# 4. EXPLICIT SAFE GLOBALS WHITELIST
+# Even with the monkeypatch, we add these to be double-safe against internal calls
 try:
-    from omegaconf.listconfig import ListConfig
-    from omegaconf.dictconfig import DictConfig
-    try:
-        from omegaconf.base import ContainerMetadata, Metadata
-    except ImportError:
-        ContainerMetadata = None
-        Metadata = None
-
     safe_list = [
-        typing.Any, # Specific fix for your error
-        ListConfig, 
-        DictConfig,
+        typing.Any,
         torch.nn.modules.container.ModuleList,
         np.dtype,
     ]
-    if ContainerMetadata: safe_list.append(ContainerMetadata)
-    if Metadata: safe_list.append(Metadata)
     
+    # NumPy internals
     if hasattr(np, '_core') and hasattr(np._core, 'multiarray'):
         safe_list.append(np._core.multiarray.scalar)
     elif hasattr(np, 'core') and hasattr(np.core, 'multiarray'):
         safe_list.append(np.core.multiarray.scalar)
 
+    # OmegaConf (Crucial for Pyannote/WhisperX config loading)
+    try:
+        from omegaconf.listconfig import ListConfig
+        from omegaconf.dictconfig import DictConfig
+        from omegaconf.base import ContainerMetadata, Metadata, Node
+        safe_list.extend([ListConfig, DictConfig, ContainerMetadata, Metadata, Node])
+    except ImportError:
+        print("Warning: Could not import omegaconf for whitelisting.")
+
+    # Pyannote internals (if available)
+    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
+        safe_list.extend([Specifications, Problem, Resolution, Model, SpeakerDiarization])
+    except ImportError:
+        pass
+
     torch.serialization.add_safe_globals(safe_list)
 except Exception as e:
     print(f"Safe Globals Warning: {e}")
@@ -73,10 +89,6 @@ ENV_GEMINI_KEY = os.environ.get("GEMINI_API_KEY", "")
 ACTIVE_HF_TOKEN = ENV_HF_TOKEN if ENV_HF_TOKEN else HARDCODED_HF_TOKEN
 ACTIVE_GEMINI_KEY = ENV_GEMINI_KEY if ENV_GEMINI_KEY else HARDCODED_GEMINI_KEY
 
-# Fix OMP Threads
-if os.environ.get("OMP_NUM_THREADS", "").endswith("m"):
-    os.environ["OMP_NUM_THREADS"] = "1"
-
 def format_timecode(seconds, fps=25):
     """Converts seconds to HH:MM:SS:FF."""
     td = timedelta(seconds=seconds)
@@ -150,6 +162,20 @@ with st.sidebar:
     
     st.header("Model Settings")
     model_size = st.selectbox("Whisper Model", ["large-v2", "medium", "base"], index=0)
+    
+    # --- Language Option ---
+    language_map = {
+        "Auto-Detect": None,
+        "English": "en",
+        "Spanish": "es",
+        "French": "fr",
+        "German": "de",
+        "Italian": "it",
+        "Portuguese": "pt"
+    }
+    selected_lang_label = st.selectbox("Audio Language", list(language_map.keys()), index=1)
+    target_language = language_map[selected_lang_label]
+    
     num_speakers = st.number_input("Speakers (0=Auto)", min_value=0, value=0)
     
     st.divider()
@@ -180,7 +206,7 @@ if uploaded_file:
                     try:
                         device = "cuda" if torch.cuda.is_available() else "cpu"
                         if device == "cpu":
-                            st.warning("⚠️ No GPU detected. WhisperX will be slow.")
+                            st.warning("⚠️ No GPU detected. WhisperX will be very slow.")
                         else:
                             st.write(f"🚀 **Loading WhisperX on {device}...**")
                         
@@ -190,8 +216,9 @@ if uploaded_file:
                         
                         st.write("📝 **Transcribing...**")
                         audio = whisperx.load_audio("temp_audio.wav")
-                        result = model.transcribe(audio, batch_size=16)
+                        result = model.transcribe(audio, batch_size=16, language=target_language)
                         
+                        # Memory cleanup
                         del model
                         gc.collect()
                         torch.cuda.empty_cache()
@@ -207,13 +234,22 @@ if uploaded_file:
                         
                         # 3. Diarize
                         st.write("🗣️ **Diarizing Speakers...**")
-                        # The monkeypatch above should protect this call
-                        diarize_model = whisperx.DiarizationPipeline(use_auth_token=ACTIVE_HF_TOKEN, device=device)
+                        # Pass token for gated diarization models
+                        # Try to bypass the torch security default if necessary
+                        try:
+                            diarize_model = whisperx.DiarizationPipeline(use_auth_token=ACTIVE_HF_TOKEN, device=device)
+                        except Exception as e:
+                            # Catch generic loading errors and try to print detail or retry
+                            if "Weights only load failed" in str(e) or "Unsupported global" in str(e):
+                                st.warning("⚠️ Security restriction encountered. Re-attempting load with implicit overrides.")
+                                diarize_model = whisperx.DiarizationPipeline(use_auth_token=ACTIVE_HF_TOKEN, device=device)
+                            else:
+                                raise e
                         
                         diarize_kwargs = {}
                         if num_speakers > 0:
                             diarize_kwargs = {"min_speakers": num_speakers, "max_speakers": num_speakers}
-                            
+                        
                         diarize_segments = diarize_model(audio, **diarize_kwargs)
                         
                         # 4. Final Merge