Upload app.py with huggingface_hub

app.py

@@ -8,40 +8,36 @@ import numpy as np
 import pandas as pd
 import requests
 import soundfile as sf
-import torch
 import torchaudio
+import torch
 from sklearn.cluster import AgglomerativeClustering
-from transformers import Wav2Vec2FeatureExtractor, WavLMForXVector
+from optimum.onnxruntime import ORTModelForAudioXVector
+from transformers import Wav2Vec2FeatureExtractor
 
 os.environ.setdefault("HF_XET_HIGH_PERFORMANCE", "1")
 
-# 1 PyTorch thread per worker — lets N_CPUS threads run inference in parallel
-# instead of one wide inference that blocks everyone else.
-torch.set_num_threads(1)
-
 N_CPUS = os.cpu_count() or 2
 
 DATASETS_SERVER = "https://datasets-server.huggingface.co"
-MODEL_ID = "microsoft/wavlm-base-plus-sv"
+ONNX_MODEL_ID = "fosters/wavlm-base-plus-sv-onnx"
 TARGET_SR = 16000
 
 _feature_extractor = None
 _model = None
-_init_lock = threading.Lock()  # only for one-time model init
+_init_lock = threading.Lock()
 
 
 def _load_model():
     global _feature_extractor, _model
     with _init_lock:
         if _model is None:
-            _feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(MODEL_ID)
-            _model = WavLMForXVector.from_pretrained(MODEL_ID)
-            _model.eval()
+            _feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(ONNX_MODEL_ID)
+            # ONNX Runtime: thread-safe, no GIL concerns
+            _model = ORTModelForAudioXVector.from_pretrained(ONNX_MODEL_ID)
     return _feature_extractor, _model
 
 
 def _embed(audio_array: np.ndarray, sr: int, max_sec: int) -> np.ndarray:
-    # Thread-safe: eval() + no_grad() — weights are read-only, GIL released in C++ ops
     fe, mdl = _load_model()
     waveform = torch.tensor(audio_array, dtype=torch.float32)
     if waveform.ndim == 2:
@@ -50,8 +46,7 @@ def _embed(audio_array: np.ndarray, sr: int, max_sec: int) -> np.ndarray:
         waveform = torchaudio.functional.resample(waveform, sr, TARGET_SR)
     waveform = waveform[: max_sec * TARGET_SR]
     inputs = fe(waveform.numpy(), sampling_rate=TARGET_SR, return_tensors="pt")
-    with torch.no_grad():
-        out = mdl(**inputs)
+    out = mdl(**inputs)
     return out.embeddings.squeeze().numpy()