From 6ec30d850595509ab417424b14f3f631c3f7d5d0 Mon Sep 17 00:00:00 2001 From: Cairn Agent Date: Mon, 20 Jul 2026 16:19:46 -0700 Subject: [PATCH 1/4] [Build] harden model weight acquisition and loading --- src/deeprhythm/batch_infer.py | 2 +- src/deeprhythm/model/predictor.py | 6 +- src/deeprhythm/utils.py | 114 ++++++++++++++++++++++++------ 3 files changed, 98 insertions(+), 24 deletions(-) diff --git a/src/deeprhythm/batch_infer.py b/src/deeprhythm/batch_infer.py index 8d3edb7..ba3a9aa 100644 --- a/src/deeprhythm/batch_infer.py +++ b/src/deeprhythm/batch_infer.py @@ -120,7 +120,7 @@ def consume_and_process( if not quiet: print('made kernels') model = DeepRhythmModel() - model.load_state_dict(torch.load(get_weights(quiet=quiet), map_location=torch.device(device), weights_only=False)) + model.load_state_dict(load_weights(torch.device(device), quiet=quiet)) model = model.to(device=device) model.eval() if not quiet: diff --git a/src/deeprhythm/model/predictor.py b/src/deeprhythm/model/predictor.py index c616895..3d85ec7 100644 --- a/src/deeprhythm/model/predictor.py +++ b/src/deeprhythm/model/predictor.py @@ -10,7 +10,7 @@ from deeprhythm.batch_infer import get_audio_files from deeprhythm.batch_infer import main as batch_infer_main from deeprhythm.model.frame_cnn import DeepRhythmModel -from deeprhythm.utils import class_to_bpm, get_device, get_weights, load_and_split_audio, split_audio +from deeprhythm.utils import class_to_bpm, get_device, load_and_split_audio, load_weights, split_audio class DeepRhythmPredictor: @@ -22,15 +22,15 @@ def __init__(self, device: Optional[str] = None, quiet: bool = False): device: Device to run inference on ('cuda' or 'cpu') quiet: Whether to suppress progress messages """ - self.model_path = get_weights(quiet=quiet) self.device = torch.device(device if device else get_device()) + self.quiet = quiet self.model = self._load_model() self.specs = self._make_kernels() def _load_model(self) -> DeepRhythmModel: """Load and initialize the model.""" model = DeepRhythmModel() - model.load_state_dict(torch.load(self.model_path, map_location=self.device, weights_only=False)) + model.load_state_dict(load_weights(self.device, quiet=self.quiet)) model = model.to(device=self.device) model.eval() return model diff --git a/src/deeprhythm/utils.py b/src/deeprhythm/utils.py index 9d425f5..40aa480 100644 --- a/src/deeprhythm/utils.py +++ b/src/deeprhythm/utils.py @@ -1,10 +1,22 @@ +import hashlib import os +import tempfile import librosa import requests import torch -model_url = 'https://github.com/bleugreen/deeprhythm/raw/main/weights/' +MODEL_WEIGHTS_FILENAME = "deeprhythm-0.7.pth" +MODEL_WEIGHTS_URL = ( + "https://raw.githubusercontent.com/bleugreen/deeprhythm/" + "c3548590d12f5d97ccd39dee27e9591e522c05e2/weights/deeprhythm-0.7.pth" +) +MODEL_WEIGHTS_SHA256 = "c7cc8cc0425929cd2bf695474d7ec1fd63ed0d0a4a68f361d4e4b57bd9b3d9c4" +MODEL_WEIGHTS_SIZE = 5_443_228 +MODEL_DOWNLOAD_CONNECT_TIMEOUT = 10 +MODEL_DOWNLOAD_READ_TIMEOUT = 60 +MODEL_DOWNLOAD_TIMEOUT = (MODEL_DOWNLOAD_CONNECT_TIMEOUT, MODEL_DOWNLOAD_READ_TIMEOUT) +MODEL_DOWNLOAD_CHUNK_SIZE = 64 * 1024 class AudioTooShortError(ValueError): @@ -12,6 +24,23 @@ class AudioTooShortError(ValueError): pass +class ModelWeightsError(RuntimeError): + """Raised when verified model weights cannot be acquired.""" + + +def _weights_are_valid(path): + try: + if os.path.getsize(path) != MODEL_WEIGHTS_SIZE: + return False + digest = hashlib.sha256() + with open(path, "rb") as weights_file: + for chunk in iter(lambda: weights_file.read(MODEL_DOWNLOAD_CHUNK_SIZE), b""): + digest.update(chunk) + return digest.hexdigest() == MODEL_WEIGHTS_SHA256 + except OSError: + return False + + class AudioLoadError(IOError): """Raised when audio file cannot be loaded.""" pass @@ -26,30 +55,75 @@ def get_device(): return 'cpu' -def get_weights(filename="deeprhythm-0.7.pth", quiet=False): +def get_weights(quiet=False): home_dir = os.path.expanduser("~") model_dir = os.path.join(home_dir, ".local", "share", "deeprhythm") - if not os.path.exists(model_dir): - os.makedirs(model_dir, exist_ok=True) - model_path = os.path.join(model_dir, filename) + os.makedirs(model_dir, exist_ok=True) + model_path = os.path.join(model_dir, MODEL_WEIGHTS_FILENAME) - if not os.path.isfile(model_path): - print("Downloading model weights...") - try: - r = requests.get(model_url + filename, allow_redirects=True) - if r.status_code == 200: - with open(model_path, 'wb') as f: - f.write(r.content) - print("Model weights downloaded successfully.") - else: - print(f"Failed to download model weights. HTTP Error: {r.status_code}") - except Exception as e: - print(f"An error occurred during the download: {e}") - else: + if _weights_are_valid(model_path): if not quiet: - print("Model weights already exist.") + print("Model weights already exist and are verified.") + return model_path - return model_path + temporary_path = None + try: + print("Downloading model weights...") + with requests.get( + MODEL_WEIGHTS_URL, + stream=True, + timeout=MODEL_DOWNLOAD_TIMEOUT, + ) as response: + response.raise_for_status() + digest = hashlib.sha256() + downloaded_size = 0 + with tempfile.NamedTemporaryFile( + mode="wb", + dir=model_dir, + prefix=f".{MODEL_WEIGHTS_FILENAME}.", + suffix=".tmp", + delete=False, + ) as temporary_file: + temporary_path = temporary_file.name + for chunk in response.iter_content(chunk_size=MODEL_DOWNLOAD_CHUNK_SIZE): + if not chunk: + continue + downloaded_size += len(chunk) + if downloaded_size > MODEL_WEIGHTS_SIZE: + raise ModelWeightsError( + f"Model weights exceed the expected size of {MODEL_WEIGHTS_SIZE} bytes" + ) + digest.update(chunk) + temporary_file.write(chunk) + + if downloaded_size != MODEL_WEIGHTS_SIZE: + raise ModelWeightsError( + f"Model weights have size {downloaded_size} bytes; expected {MODEL_WEIGHTS_SIZE}" + ) + if digest.hexdigest() != MODEL_WEIGHTS_SHA256: + raise ModelWeightsError("Model weights failed SHA-256 verification") + + os.replace(temporary_path, model_path) + temporary_path = None + if not quiet: + print("Model weights downloaded and verified successfully.") + return model_path + except requests.RequestException as exc: + raise ModelWeightsError(f"Failed to download model weights: {exc}") from exc + except OSError as exc: + raise ModelWeightsError(f"Failed to install model weights: {exc}") from exc + finally: + if temporary_path is not None: + try: + os.unlink(temporary_path) + except FileNotFoundError: + pass + + +def load_weights(device, quiet=False): + """Acquire verified weights and safely deserialize their state dictionary.""" + path = get_weights(quiet=quiet) + return torch.load(path, map_location=device, weights_only=True) def split_audio(audio, sr, clip_length=8, share_mem=False): From 73a90b0d4b6d626ba6fc678a8f1d15bfc82de8ae Mon Sep 17 00:00:00 2001 From: bleugreen Date: Mon, 20 Jul 2026 16:19:47 -0700 Subject: [PATCH 2/4] [Build] add model weight security regression tests --- src/deeprhythm/batch_infer.py | 2 +- src/deeprhythm/utils.py | 3 +- tests/test_utils.py | 147 ++++++++++++++++++++++++++++++++++ 3 files changed, 150 insertions(+), 2 deletions(-) diff --git a/src/deeprhythm/batch_infer.py b/src/deeprhythm/batch_infer.py index ba3a9aa..f5614a1 100644 --- a/src/deeprhythm/batch_infer.py +++ b/src/deeprhythm/batch_infer.py @@ -14,8 +14,8 @@ AudioTooShortError, class_to_bpm, get_device, - get_weights, load_and_split_audio, + load_weights, ) NUM_WORKERS = 8 diff --git a/src/deeprhythm/utils.py b/src/deeprhythm/utils.py index 40aa480..c2ecefa 100644 --- a/src/deeprhythm/utils.py +++ b/src/deeprhythm/utils.py @@ -68,7 +68,8 @@ def get_weights(quiet=False): temporary_path = None try: - print("Downloading model weights...") + if not quiet: + print("Downloading model weights...") with requests.get( MODEL_WEIGHTS_URL, stream=True, diff --git a/tests/test_utils.py b/tests/test_utils.py index 685ad09..637d8ab 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,4 +1,12 @@ +import hashlib +from pathlib import Path +from unittest.mock import Mock + import numpy as np +import pytest +import requests + +import deeprhythm.utils as utils from deeprhythm.utils import ( AudioLoadError, @@ -6,9 +14,148 @@ bpm_to_class, class_to_bpm, get_device, + get_weights, + load_weights, split_audio, ) + +class StreamedResponse: + def __init__(self, chunks=(), error=None): + self.chunks = chunks + self.error = error + + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + def raise_for_status(self): + if self.error: + raise self.error + + def iter_content(self, chunk_size): + assert chunk_size == utils.MODEL_DOWNLOAD_CHUNK_SIZE + yield from self.chunks + + +@pytest.fixture +def model_artifact(monkeypatch, tmp_path): + payload = b"verified model bytes" + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.setattr(utils, "MODEL_WEIGHTS_SIZE", len(payload)) + monkeypatch.setattr(utils, "MODEL_WEIGHTS_SHA256", hashlib.sha256(payload).hexdigest()) + cache_dir = tmp_path / ".local" / "share" / "deeprhythm" + destination = cache_dir / utils.MODEL_WEIGHTS_FILENAME + return payload, cache_dir, destination + + +def test_get_weights_accepts_valid_cache_without_network(model_artifact, monkeypatch): + payload, cache_dir, destination = model_artifact + cache_dir.mkdir(parents=True) + destination.write_bytes(payload) + request = Mock(side_effect=AssertionError("network must not be used")) + monkeypatch.setattr(utils.requests, "get", request) + + assert get_weights(quiet=True) == str(destination) + request.assert_not_called() + + +def test_get_weights_downloads_from_pinned_url_and_atomically_replaces_corrupt_cache( + model_artifact, monkeypatch +): + payload, cache_dir, destination = model_artifact + cache_dir.mkdir(parents=True) + destination.write_bytes(b"corrupt") + request = Mock(return_value=StreamedResponse([payload[:5], payload[5:]])) + monkeypatch.setattr(utils.requests, "get", request) + real_replace = utils.os.replace + replace = Mock(side_effect=real_replace) + monkeypatch.setattr(utils.os, "replace", replace) + + assert get_weights(quiet=True) == str(destination) + request.assert_called_once_with( + utils.MODEL_WEIGHTS_URL, + stream=True, + timeout=(utils.MODEL_DOWNLOAD_CONNECT_TIMEOUT, utils.MODEL_DOWNLOAD_READ_TIMEOUT), + ) + source, target = replace.call_args.args + assert target == str(destination) + assert destination.read_bytes() == payload + assert not Path(source).exists() + assert list(cache_dir.glob("*.tmp")) == [] + + +@pytest.mark.parametrize( + ("response", "expected_message"), + [ + (StreamedResponse(error=requests.HTTPError("500 Server Error")), "Failed to download"), + (StreamedResponse([b"partial"], requests.Timeout("timed out")), "Failed to download"), + ], +) +def test_get_weights_network_failures_do_not_publish( + response, expected_message, model_artifact, monkeypatch +): + _, cache_dir, destination = model_artifact + monkeypatch.setattr(utils.requests, "get", Mock(return_value=response)) + + with pytest.raises(utils.ModelWeightsError, match=expected_message): + get_weights(quiet=True) + + assert not destination.exists() + assert list(cache_dir.glob("*.tmp")) == [] + + +def test_get_weights_interrupted_stream_preserves_existing_cache(model_artifact, monkeypatch): + payload, cache_dir, destination = model_artifact + cache_dir.mkdir(parents=True) + destination.write_bytes(b"existing corrupt bytes") + + class InterruptedResponse(StreamedResponse): + def iter_content(self, chunk_size): + yield payload[:5] + raise requests.ConnectionError("stream interrupted") + + monkeypatch.setattr(utils.requests, "get", Mock(return_value=InterruptedResponse())) + + with pytest.raises(utils.ModelWeightsError, match="stream interrupted"): + get_weights(quiet=True) + + assert destination.read_bytes() == b"existing corrupt bytes" + assert list(cache_dir.glob("*.tmp")) == [] + + +@pytest.mark.parametrize( + ("download", "expected_message"), + [ + (b"verified model bytes!", "exceed the expected size"), + (b"short", "have size"), + (b"different model bytes", "SHA-256"), + ], +) +def test_get_weights_rejects_invalid_downloads(download, expected_message, model_artifact, monkeypatch): + _, cache_dir, destination = model_artifact + monkeypatch.setattr(utils.requests, "get", Mock(return_value=StreamedResponse([download]))) + + with pytest.raises(utils.ModelWeightsError, match=expected_message): + get_weights(quiet=True) + + assert not destination.exists() + assert list(cache_dir.glob("*.tmp")) == [] + + +def test_load_weights_uses_safe_torch_deserialization(monkeypatch): + state_dict = {"layer": "weights"} + monkeypatch.setattr(utils, "get_weights", Mock(return_value="/verified/model.pth")) + torch_load = Mock(return_value=state_dict) + monkeypatch.setattr(utils.torch, "load", torch_load) + + assert load_weights("cpu", quiet=True) == state_dict + torch_load.assert_called_once_with( + "/verified/model.pth", map_location="cpu", weights_only=True + ) + # --------------------------------------------------------------------------- # BPM <-> class conversion # --------------------------------------------------------------------------- From 656653b5b2503113e473785460beeef2073bc7cc Mon Sep 17 00:00:00 2001 From: bleugreen Date: Mon, 20 Jul 2026 16:20:15 -0700 Subject: [PATCH 3/4] fix test import ordering --- tests/test_utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_utils.py b/tests/test_utils.py index 637d8ab..f822c30 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -7,7 +7,6 @@ import requests import deeprhythm.utils as utils - from deeprhythm.utils import ( AudioLoadError, AudioTooShortError, From 630f126c9bec48255e7d82cc21267a36e9390c63 Mon Sep 17 00:00:00 2001 From: bleugreen Date: Mon, 20 Jul 2026 16:20:54 -0700 Subject: [PATCH 4/4] [Build] correct digest mismatch fixture size --- tests/test_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_utils.py b/tests/test_utils.py index f822c30..13476f9 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -130,7 +130,7 @@ def iter_content(self, chunk_size): [ (b"verified model bytes!", "exceed the expected size"), (b"short", "have size"), - (b"different model bytes", "SHA-256"), + (b"tampered model bytes", "SHA-256"), ], ) def test_get_weights_rejects_invalid_downloads(download, expected_message, model_artifact, monkeypatch):