Palash DebnathandClaude Opus 4.8
f7d34a1433
On torch>=2.6, `Pipeline.from_pretrained("pyannote/speaker-diarization-3.1")`
fails with "Weights only load failed ... Unsupported global: GLOBAL
torch.torch_version.TorchVersion" — PyTorch 2.6 flipped torch.load's default to
weights_only=True and its secure unpickler rejects the checkpoint's metadata
globals. This broke diarization on torch>=2.6 even when the license IS accepted
(reported on v0.3.4, RTX 4070 Ti, license accepted).
The WhisperX VAD load already solved this via
`WhisperXBackend._allow_vad_pickle_globals()` (allowlists TorchVersion,
omegaconf nodes, pyannote metadata, builtins, numpy, …). `get_diarization_pipeline`
just never called it. Reuse it before the diarization load — idempotent,
per-process, verified to register TorchVersion on torch 2.8.
Graceful fallback (silence-gap heuristic) is preserved if anything still fails.
Tests: tests/test_diarization_weights_only.py (allowlist runs before load;
no-token short-circuit). Existing diarization classification tests still pass.
Cross-platform (the torch 2.6 weights_only change affects all platforms).
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
66 lines
2.1 KiB
Python
66 lines
2.1 KiB
Python
"""Diarization must register torch safe-globals before loading (issue #270).
|
|
|
|
PyTorch 2.6+ defaults `torch.load` to `weights_only=True`, whose secure
|
|
unpickler rejects the pyannote checkpoint's metadata globals
|
|
(`torch_version.TorchVersion`, omegaconf nodes, …). `get_diarization_pipeline`
|
|
must register the same allowlist the WhisperX VAD load uses, before calling
|
|
`Pipeline.from_pretrained`, or diarization breaks even with the license
|
|
accepted.
|
|
"""
|
|
import sys
|
|
import types
|
|
|
|
import pytest
|
|
|
|
|
|
@pytest.fixture
|
|
def reset_diar(monkeypatch):
|
|
import services.model_manager as mm
|
|
monkeypatch.setattr(mm, "_diar_pipeline", None, raising=False)
|
|
yield mm
|
|
monkeypatch.setattr(mm, "_diar_pipeline", None, raising=False)
|
|
|
|
|
|
def test_loads_pyannote_after_registering_safe_globals(reset_diar, monkeypatch):
|
|
mm = reset_diar
|
|
order = []
|
|
|
|
# Token present (App source).
|
|
monkeypatch.setattr(
|
|
"services.token_resolver.resolve",
|
|
lambda: types.SimpleNamespace(token="hf_test", source="app", user="u"),
|
|
)
|
|
|
|
# Spy on the shared allowlister; must run BEFORE from_pretrained.
|
|
from services import asr_backend as ab
|
|
monkeypatch.setattr(
|
|
ab.WhisperXBackend, "_allow_vad_pickle_globals",
|
|
staticmethod(lambda: order.append("allow")),
|
|
)
|
|
|
|
fake_pipe = object()
|
|
|
|
def _from_pretrained(*a, **k):
|
|
order.append("load")
|
|
return fake_pipe
|
|
|
|
fake_mod = types.ModuleType("pyannote.audio")
|
|
fake_mod.Pipeline = types.SimpleNamespace(from_pretrained=_from_pretrained)
|
|
monkeypatch.setitem(sys.modules, "pyannote.audio", fake_mod)
|
|
|
|
# CPU device → no .to() call on the fake pipe.
|
|
monkeypatch.setattr(mm, "get_best_device", lambda: "cpu")
|
|
|
|
result = mm.get_diarization_pipeline()
|
|
|
|
assert result is fake_pipe
|
|
assert order == ["allow", "load"], f"allowlist must precede load, got {order}"
|
|
|
|
|
|
def test_no_token_short_circuits_without_loading(reset_diar, monkeypatch):
|
|
mm = reset_diar
|
|
monkeypatch.setattr("services.token_resolver.resolve", lambda: None)
|
|
pipe, err = mm.get_diarization_pipeline(return_error=True)
|
|
assert pipe is None
|
|
assert err == mm.DIARIZATION_ERR_NO_TOKEN
|