Upload app.py with huggingface_hub

app.py

@@ -15,18 +15,24 @@ from transformers import Wav2Vec2FeatureExtractor, WavLMForXVector
 
 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"
 TARGET_SR = 16000
 
 _feature_extractor = None
 _model = None
-_embed_lock = threading.Lock()  # WavLM inference is not thread-safe
+_init_lock = threading.Lock()  # only for one-time model init
 
 
 def _load_model():
     global _feature_extractor, _model
-    with _embed_lock:
+    with _init_lock:
         if _model is None:
             _feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(MODEL_ID)
             _model = WavLMForXVector.from_pretrained(MODEL_ID)
@@ -35,6 +41,7 @@ def _load_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:
@@ -43,9 +50,8 @@ 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 _embed_lock:
-        with torch.no_grad():
-            out = mdl(**inputs)
+    with torch.no_grad():
+        out = mdl(**inputs)
     return out.embeddings.squeeze().numpy()
 
 
@@ -186,8 +192,8 @@ def identify_speakers(
     errors: list[str] = []
     done = 0
 
-    # Process all repos in parallel (downloads are I/O bound)
-    with concurrent.futures.ThreadPoolExecutor(max_workers=8) as ex:
+    # Process repos in parallel — capped at N_CPUS since embed is the bottleneck
+    with concurrent.futures.ThreadPoolExecutor(max_workers=N_CPUS) as ex:
         future_to_repo = {
             ex.submit(_process_repo, repo, int(samples_per_book), int(audio_sec), token): repo
             for repo in repos