Harden checkpoint loading: set weights_only=True on all torch.load calls

The fine-tuning resume path in scripts/train_voxcpm_finetune.py called
torch.load() without weights_only=True for the LoRA checkpoint, full-model
checkpoint, optimizer, and scheduler. Loading an attacker-supplied checkpoint
directory would therefore execute arbitrary code during unpickling.

Every inference-time loader (model/voxcpm.py, model/voxcpm2.py, LoRA loading)
already passes weights_only=True, and the project ships a test asserting LoRA
loading rejects malicious pickle payloads. This closes the remaining gap so the
resume path matches that posture.

Also:
- app.py: add a --host flag so the Gradio server can bind to 127.0.0.1 instead
  of being hardwired to 0.0.0.0 (default unchanged to preserve current behavior).
  The UI exposes an unauthenticated api_name="generate" endpoint.
- tests/test_torch_load_safety.py: AST guard asserting every torch.load across
  src/, scripts/, app.py and lora_ft_webui.py sets weights_only=True, plus a
  behavioral check that a malicious pickle is blocked.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

app.py

@@ -529,6 +529,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,
@@ -536,4 +543,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,
+    )

scripts/train_voxcpm_finetune.py

@@ -682,7 +682,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)
@@ -700,7 +700,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)
@@ -710,14 +710,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)
 

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"