fix: refuse untrusted globals when loading model checkpoints (v0.47.2)

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:

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

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()

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}.")
 

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):

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 <andrew@beveridge.uk>"]
 license = "MIT"

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))

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"]
 
 

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