From ea46259160fd7a9ba4a3ba73be2533f7a1d3a291 Mon Sep 17 00:00:00 2001 From: Andrew Beveridge Date: Thu, 29 Oct 2026 02:06:40 -0400 Subject: [PATCH] fix: refuse untrusted globals when loading model checkpoints (v0.47.2) Demucs package files are pickles that store the model class, and were loaded with torch.load(weights_only=False), so a crafted .th file could run code at load time. Load them through a restricted unpickler that only resolves the globals real Demucs v4 packages use (verified byte-identical against all 7 published checkpoints). Pass weights_only=True explicitly on the VR, MDXC, RoFormer and MDX mixer loads, which were only safe by default on torch >= 2.6, and drop the RoFormer legacy fallback that retried without it. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_012RYLVTA9CYSnmncNv3GRHJ --- .../separator/architectures/mdxc_separator.py | 2 +- .../separator/architectures/vr_separator.py | 2 +- .../separator/roformer/roformer_loader.py | 8 +-- .../separator/uvr_lib_v5/demucs/states.py | 48 +++++++++++-- .../separator/uvr_lib_v5/mdxnet.py | 2 +- pyproject.toml | 2 +- tests/unit/test_demucs_safe_load.py | 67 +++++++++++++++++++ tests/unit/test_directml.py | 4 +- tests/unit/test_model_reuse.py | 2 +- 9 files changed, 121 insertions(+), 16 deletions(-) create mode 100644 tests/unit/test_demucs_safe_load.py diff --git a/audio_separator/separator/architectures/mdxc_separator.py b/audio_separator/separator/architectures/mdxc_separator.py index c21d1a53..d03f2be2 100644 --- a/audio_separator/separator/architectures/mdxc_separator.py +++ b/audio_separator/separator/architectures/mdxc_separator.py @@ -196,7 +196,7 @@ def load_model(self): self.logger.debug("Loading model onto cpu") # For some reason loading the state onto a hardware accelerated devices causes issues, # so we load it onto CPU first then move it to the device - self.model_run.load_state_dict(torch.load(self.model_path, map_location="cpu")) + self.model_run.load_state_dict(torch.load(self.model_path, map_location="cpu", weights_only=True)) self.model_run.to(self.torch_device).eval() except RuntimeError as e: diff --git a/audio_separator/separator/architectures/vr_separator.py b/audio_separator/separator/architectures/vr_separator.py index 3c5c83eb..1de806e5 100644 --- a/audio_separator/separator/architectures/vr_separator.py +++ b/audio_separator/separator/architectures/vr_separator.py @@ -132,7 +132,7 @@ def _ensure_model_loaded(self, nn_arch_size): # Publish the reusable module only after every loading step succeeds. # Otherwise a retry could mistake a partial model for a ready one. - model_run.load_state_dict(torch.load(self.model_path, map_location="cpu")) + model_run.load_state_dict(torch.load(self.model_path, map_location="cpu", weights_only=True)) model_run.to(self.torch_device) self.model_run = model_run self.is_vr_51_model = is_vr_51_model diff --git a/audio_separator/separator/roformer/roformer_loader.py b/audio_separator/separator/roformer/roformer_loader.py index 8994ec7f..954b699c 100644 --- a/audio_separator/separator/roformer/roformer_loader.py +++ b/audio_separator/separator/roformer/roformer_loader.py @@ -116,7 +116,7 @@ def _load_with_new_implementation(self, # everywhere, but is gated to DML so all other devices keep # their exact existing behavior. (Issue #292) map_location = "cpu" if _is_dml_device(device) else device - state_dict = torch.load(model_path, map_location=map_location) + state_dict = torch.load(model_path, map_location=map_location, weights_only=True) if isinstance(state_dict, dict) and 'state_dict' in state_dict: model.load_state_dict(state_dict['state_dict']) elif isinstance(state_dict, dict) and 'model' in state_dict: @@ -247,11 +247,7 @@ def _load_with_legacy_implementation(self, raise ValueError(f"Unknown Roformer model type in legacy configuration: {model_type}") # Load checkpoint as raw state dict (legacy behavior) - try: - checkpoint = torch.load(model_path, map_location='cpu', weights_only=True) - except TypeError: - # For older torch versions without weights_only - checkpoint = torch.load(model_path, map_location='cpu') + checkpoint = torch.load(model_path, map_location='cpu', weights_only=True) model.load_state_dict(checkpoint) model.to(device).eval() diff --git a/audio_separator/separator/uvr_lib_v5/demucs/states.py b/audio_separator/separator/uvr_lib_v5/demucs/states.py index 9bf1a8af..8a9149d1 100644 --- a/audio_separator/separator/uvr_lib_v5/demucs/states.py +++ b/audio_separator/separator/uvr_lib_v5/demucs/states.py @@ -12,6 +12,7 @@ import hashlib import inspect import io +import pickle from pathlib import Path import warnings @@ -31,16 +32,55 @@ def get_quantizer(model, args, optimizer=None): return quantizer +# Every global a Demucs v4 package (.th) references, as (module, name). Model files are +# pickles, so resolving anything outside this list could run arbitrary code at load time. +# torch storage types never reach find_class: torch's UnpicklerWrapper handles them first. +_ALLOWED_PACKAGE_GLOBALS = frozenset( + { + ("collections", "OrderedDict"), + ("torch._utils", "_rebuild_tensor_v2"), + ("torch._utils", "_rebuild_parameter"), + ("_codecs", "encode"), + ("fractions", "Fraction"), + ("numpy", "dtype"), + ("numpy.core.multiarray", "scalar"), + ("numpy._core.multiarray", "scalar"), + ("demucs.demucs", "Demucs"), + ("demucs.hdemucs", "HDemucs"), + ("demucs.htdemucs", "HTDemucs"), + } +) + + +class _RestrictedUnpickler(pickle.Unpickler): + def find_class(self, module, name): + if (module, name) in _ALLOWED_PACKAGE_GLOBALS or (module == "numpy.dtypes" and name.endswith("DType")): + return super().find_class(module, name) + raise pickle.UnpicklingError(f"Refusing to load Demucs model: global '{module}.{name}' is not allowed") + + +class _RestrictedPickleModule: + """pickle-compatible module for torch.load(pickle_module=...) that only resolves allowed globals.""" + + __name__ = "audio_separator_restricted_pickle" + Unpickler = _RestrictedUnpickler + load = staticmethod(lambda f, **kwargs: _RestrictedUnpickler(f, **kwargs).load()) + + +def _load_package(path): + # Package files store the model class itself, which torch's weights_only loader can't + # express portably across the torch/numpy versions we support, so restrict the + # unpickler instead. weights_only must be False for torch to use pickle_module. + return torch.load(path, "cpu", pickle_module=_RestrictedPickleModule, weights_only=False) + + def load_model(path_or_package, strict=False): """Load a model from the given serialized model, either given as a dict (already loaded) or a path to a file on disk.""" if isinstance(path_or_package, dict): package = path_or_package elif isinstance(path_or_package, (str, Path)): - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - path = path_or_package - package = torch.load(path, "cpu", weights_only=False) + package = _load_package(path_or_package) else: raise ValueError(f"Invalid type for {path_or_package}.") diff --git a/audio_separator/separator/uvr_lib_v5/mdxnet.py b/audio_separator/separator/uvr_lib_v5/mdxnet.py index 3293c89f..f8b5a4e8 100644 --- a/audio_separator/separator/uvr_lib_v5/mdxnet.py +++ b/audio_separator/separator/uvr_lib_v5/mdxnet.py @@ -127,7 +127,7 @@ def __init__(self, device, mixer_path): self.linear = nn.Linear((dim_s+1)*2, dim_s*2, bias=False) self.load_state_dict( - torch.load(mixer_path, map_location=device) + torch.load(mixer_path, map_location=device, weights_only=True) ) def forward(self, x): diff --git a/pyproject.toml b/pyproject.toml index 1f50d97f..b1169c4c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,7 +52,7 @@ dml = ["onnxruntime-directml>=1.17", "torch_directml"] [tool.poetry] requires-poetry = ">=2.0.0" -version = "0.47.1" +version = "0.47.2" description = "Easy to use audio stem separation, using various models from UVR trained primarily by @Anjok07" authors = ["Andrew Beveridge "] license = "MIT" diff --git a/tests/unit/test_demucs_safe_load.py b/tests/unit/test_demucs_safe_load.py new file mode 100644 index 00000000..2e5e6fa6 --- /dev/null +++ b/tests/unit/test_demucs_safe_load.py @@ -0,0 +1,67 @@ +"""Demucs package files are pickles; loading one must never run code it carries.""" + +import importlib +import os +from fractions import Fraction +from pathlib import Path + +import numpy as np +import pytest +import torch + +UVR_LIB_PATH = Path(__file__).resolve().parents[2] / "audio_separator" / "separator" / "uvr_lib_v5" + + +@pytest.fixture +def demucs_modules(monkeypatch): + # Checkpoints reference the top-level `demucs.*` names, as DemucsSeparator sets up. + monkeypatch.syspath_prepend(str(UVR_LIB_PATH)) + return importlib.import_module("demucs.states"), importlib.import_module("demucs.hdemucs") + + +class _Payload: + def __init__(self, marker): + self.marker = marker + + def __reduce__(self): + return (os.system, (f"touch {self.marker}",)) + + +def _save_package(path, klass, kwargs, state): + torch.save({"klass": klass, "args": [], "kwargs": kwargs, "state": state}, path) + + +def test_malicious_package_is_refused_without_running(demucs_modules, tmp_path): + states, _ = demucs_modules + marker = tmp_path / "pwned" + model_path = tmp_path / "evil.th" + torch.save({"klass": _Payload(str(marker)), "args": [], "kwargs": {}, "state": {}}, model_path) + + with pytest.raises(Exception, match="not allowed"): + states.load_model(model_path) + + assert not marker.exists() + + +def test_disallowed_class_reference_is_refused(demucs_modules, tmp_path): + states, _ = demucs_modules + model_path = tmp_path / "other_class.th" + _save_package(model_path, torch.nn.Linear, {"in_features": 2, "out_features": 2}, {}) + + with pytest.raises(Exception, match="torch.nn.modules.linear.Linear"): + states.load_model(model_path) + + +def test_legitimate_package_round_trips(demucs_modules, tmp_path): + states, hdemucs = demucs_modules + kwargs = {"sources": ["drums", "bass"], "channels": 4, "depth": 2, "segment": Fraction(39, 5), "samplerate": np.int64(44100)} + source = hdemucs.HDemucs(**kwargs) + model_path = tmp_path / "tiny.th" + _save_package(model_path, hdemucs.HDemucs, kwargs, {k: v.half() for k, v in source.state_dict().items()}) + + model = states.load_model(model_path) + + assert isinstance(model, hdemucs.HDemucs) + assert model.sources == ["drums", "bass"] + for key, value in source.state_dict().items(): + assert torch.equal(model.state_dict()[key], value.half().to(value.dtype)) diff --git a/tests/unit/test_directml.py b/tests/unit/test_directml.py index 8a91ed4b..9e5751e3 100644 --- a/tests/unit/test_directml.py +++ b/tests/unit/test_directml.py @@ -76,8 +76,9 @@ def _load_via_new_implementation(device): model = MagicMock(name="model") seen = {} - def fake_torch_load(path, map_location=None): + def fake_torch_load(path, map_location=None, weights_only=None): seen["map_location"] = map_location + seen["weights_only"] = weights_only return {} with patch.object(loader, "_create_bs_roformer", return_value=model), \ @@ -92,6 +93,7 @@ def fake_torch_load(path, map_location=None): assert result.success model.to.assert_called_once_with(device) + assert seen["weights_only"] is True return seen["map_location"] diff --git a/tests/unit/test_model_reuse.py b/tests/unit/test_model_reuse.py index 89b048bc..8cc72769 100644 --- a/tests/unit/test_model_reuse.py +++ b/tests/unit/test_model_reuse.py @@ -304,7 +304,7 @@ def test_vr_model_loads_placeholder_once(): separator._ensure_model_loaded(31191) determine_model_capacity.assert_called_once_with(256, 31191) - load_weights.assert_called_once_with("/tmp/model.pth", map_location="cpu") + load_weights.assert_called_once_with("/tmp/model.pth", map_location="cpu", weights_only=True) loaded_model.load_state_dict.assert_called_once_with(state_dict) loaded_model.to.assert_called_once_with(torch.device("cpu")) assert separator.model_run is loaded_model