diff --git a/app.py b/app.py index 99a3c5c..1544828 100644 --- a/app.py +++ b/app.py @@ -584,6 +584,13 @@ if __name__ == "__main__": help="Local path or HuggingFace repo ID (default: openbmb/VoxCPM2)", ) parser.add_argument("--port", type=int, default=8808, help="Server port") + parser.add_argument( + "--host", + type=str, + default="0.0.0.0", + help="Bind address. Use 127.0.0.1 to restrict access to the local machine; " + "the default 0.0.0.0 exposes the unauthenticated UI/API to the network (default: 0.0.0.0)", + ) parser.add_argument( "--device", type=str, @@ -591,4 +598,9 @@ if __name__ == "__main__": help="Runtime device: auto, cpu, mps, cuda, or cuda:N (default: auto)", ) args = parser.parse_args() - run_demo(model_id=args.model_id, server_port=args.port, device=args.device) + run_demo( + model_id=args.model_id, + server_name=args.host, + server_port=args.port, + device=args.device, + ) diff --git a/scripts/train_voxcpm_finetune.py b/scripts/train_voxcpm_finetune.py index 56e2c6d..0d75be8 100644 --- a/scripts/train_voxcpm_finetune.py +++ b/scripts/train_voxcpm_finetune.py @@ -684,7 +684,7 @@ def load_checkpoint(model, optimizer, scheduler, save_dir: Path, rank: int = 0): state_dict = load_file(str(lora_weights_path)) else: - ckpt = torch.load(lora_weights_path, map_location="cpu") + ckpt = torch.load(lora_weights_path, map_location="cpu", weights_only=True) state_dict = ckpt.get("state_dict", ckpt) unwrapped.load_state_dict(state_dict, strict=False) @@ -702,7 +702,7 @@ def load_checkpoint(model, optimizer, scheduler, save_dir: Path, rank: int = 0): state_dict = load_file(str(model_path)) else: - ckpt = torch.load(model_path, map_location="cpu") + ckpt = torch.load(model_path, map_location="cpu", weights_only=True) state_dict = ckpt.get("state_dict", ckpt) unwrapped.load_state_dict(state_dict, strict=False) @@ -712,14 +712,14 @@ def load_checkpoint(model, optimizer, scheduler, save_dir: Path, rank: int = 0): # Load optimizer state optimizer_path = latest_folder / "optimizer.pth" if optimizer_path.exists(): - optimizer.load_state_dict(torch.load(optimizer_path, map_location="cpu")) + optimizer.load_state_dict(torch.load(optimizer_path, map_location="cpu", weights_only=True)) if rank == 0: print(f"Loaded optimizer state from {optimizer_path}", file=sys.stderr) # Load scheduler state scheduler_path = latest_folder / "scheduler.pth" if scheduler_path.exists(): - scheduler.load_state_dict(torch.load(scheduler_path, map_location="cpu")) + scheduler.load_state_dict(torch.load(scheduler_path, map_location="cpu", weights_only=True)) if rank == 0: print(f"Loaded scheduler state from {scheduler_path}", file=sys.stderr) diff --git a/tests/test_torch_load_safety.py b/tests/test_torch_load_safety.py new file mode 100644 index 0000000..6f84494 --- /dev/null +++ b/tests/test_torch_load_safety.py @@ -0,0 +1,85 @@ +"""Regression guard: every ``torch.load`` call must set ``weights_only=True``. + +VoxCPM deliberately loads checkpoints with ``weights_only=True`` so that a +crafted ``.ckpt``/``.pth``/``.bin`` file cannot execute arbitrary code via +pickle during unpickling (see +``tests/test_lora_checkpoint_loading.py::test_load_lora_weights_rejects_malicious_pickle_payloads``). + +The fine-tuning resume path in ``scripts/train_voxcpm_finetune.py`` originally +called ``torch.load`` without that flag, leaving an arbitrary-code-execution +gap when resuming from an attacker-supplied checkpoint directory. This test +statically asserts the flag is present on every ``torch.load`` call across the +package and scripts so the gap cannot silently reappear. +""" + +import ast +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[1] + +# Directories whose Python files load checkpoints at runtime / on resume. +SCANNED_DIRS = [REPO_ROOT / "src", REPO_ROOT / "scripts", REPO_ROOT / "app.py", REPO_ROOT / "lora_ft_webui.py"] + + +def _python_files(): + for entry in SCANNED_DIRS: + if entry.is_file() and entry.suffix == ".py": + yield entry + elif entry.is_dir(): + yield from entry.rglob("*.py") + + +def _is_torch_load(node: ast.Call) -> bool: + func = node.func + # Matches ``torch.load(...)`` and ``load(...)`` aliased from torch. + if isinstance(func, ast.Attribute) and func.attr == "load": + return isinstance(func.value, ast.Name) and func.value.id == "torch" + return False + + +def _has_weights_only_true(node: ast.Call) -> bool: + for kw in node.keywords: + if kw.arg == "weights_only": + return isinstance(kw.value, ast.Constant) and kw.value.value is True + return False + + +def test_every_torch_load_sets_weights_only_true(): + offenders = [] + checked = 0 + for path in _python_files(): + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + for node in ast.walk(tree): + if isinstance(node, ast.Call) and _is_torch_load(node): + checked += 1 + if not _has_weights_only_true(node): + offenders.append(f"{path.relative_to(REPO_ROOT)}:{node.lineno}") + + assert checked > 0, "expected to find at least one torch.load call to verify" + assert not offenders, ( + "torch.load without weights_only=True (pickle RCE risk):\n " + + "\n ".join(offenders) + ) + + +def test_torch_load_weights_only_blocks_malicious_pickle(tmp_path): + """Behavioral check that weights_only=True actually rejects a code-exec payload.""" + torch = pytest.importorskip("torch") + + marker = tmp_path / "pwned.txt" + + class Exploit: + def __reduce__(self): + import pathlib + + return (pathlib.Path.write_text, (marker, "executed\n")) + + ckpt = tmp_path / "optimizer.pth" + torch.save({"state_dict": {"w": torch.zeros(1)}, "boom": Exploit()}, ckpt) + + with pytest.raises(Exception): + torch.load(ckpt, map_location="cpu", weights_only=True) + + assert not marker.exists(), "malicious pickle executed despite weights_only=True"