Update app.py

app.py

@@ -1,54 +1,62 @@
 import os
 import numpy as np
 import torch
-import torchaudio
 
 # --- CRITICAL ENVIRONMENT FIXES ---
-# 1. Fix for Hugging Face millicore OMP_NUM_THREADS error (e.g., "7500m")
+# 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 to avoid VAD hang
+# 2. GLOBAL PYTORCH 2.6+ SECURITY BYPASS (MONKEYPATCH)
+# This forces torch.load to behave like older versions (trusting the file).
+# Essential for loading 3rd party checkpoints (WhisperX/Pyannote) that aren't updated yet.
+_original_torch_load = torch.load
+def patched_torch_load(*args, **kwargs):
+    if 'weights_only' not in kwargs:
+        kwargs['weights_only'] = False
+    return _original_torch_load(*args, **kwargs)
+torch.load = patched_torch_load
+
+# 3. Force Torchaudio Backend
+import torchaudio
 try:
     if "ffmpeg" in torchaudio.list_audio_backends():
         torchaudio.set_audio_backend("ffmpeg")
 except Exception:
     pass
 
-# 3. Monkeypatch torch.load to default weights_only=False for PyTorch 2.6+
-import torch.serialization
-original_load = torch.load
-def patched_load(*args, **kwargs):
-    if 'weights_only' not in kwargs:
-        kwargs['weights_only'] = False
-    return original_load(*args, **kwargs)
-torch.load = patched_load
-
-# --- PYTORCH 2.6+ COMPATIBILITY PATCHES ---
+# --- SAFE GLOBALS WHITELIST (Backup Strategy) ---
+# Even with the monkeypatch, we add these to be double-safe.
+safe_globals = []
 try:
+    # PyTorch internals
+    safe_globals.append(torch.nn.modules.container.ModuleList)
+    
+    # NumPy internals
+    safe_globals.append(np.dtype)
+    if hasattr(np, '_core') and hasattr(np._core, 'multiarray'):
+        safe_globals.append(np._core.multiarray.scalar)
+    elif hasattr(np, 'core') and hasattr(np.core, 'multiarray'):
+        safe_globals.append(np.core.multiarray.scalar)
+
+    # OmegaConf (Used by Pyannote config)
+    import omegaconf
     from omegaconf.listconfig import ListConfig
     from omegaconf.dictconfig import DictConfig
+    from omegaconf.base import ContainerMetadata, Metadata, Node
+    safe_globals.extend([ListConfig, DictConfig, ContainerMetadata, Metadata, Node])
+    
+    # Pyannote internals (if available at this stage)
+    # We wrap these in try/except in case pyannote isn't fully loaded yet
     try:
-        from omegaconf.base import ContainerMetadata, Metadata
+        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_globals.extend([Specifications, Problem, Resolution, Model, SpeakerDiarization])
     except ImportError:
-        ContainerMetadata = None
-        Metadata = None
+        pass
 
-    safe_list = [
-        ListConfig, 
-        DictConfig,
-        torch.nn.modules.container.ModuleList,
-        np.dtype,
-    ]
-    if ContainerMetadata: safe_list.append(ContainerMetadata)
-    if Metadata: safe_list.append(Metadata)
-    
-    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)
-
-    torch.serialization.add_safe_globals(safe_list)
+    torch.serialization.add_safe_globals(safe_globals)
 except Exception as e:
     print(f"Safe Globals Warning: {e}")
 
@@ -227,7 +235,7 @@ if uploaded_file:
                         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