fix: use weights_only=True with explicit safe globals for torch.load

Replaces weights_only=False with weights_only=True and an explicit
add_safe_globals allowlist ([argparse.Namespace, PosixPath]) in both
JoyVASA export scripts. This is the sanctioned PyTorch mitigation for
the trailofbits pickles-in-pytorch semgrep pattern: only the known-safe
non-default classes are allowlisted; all others are rejected by the
restricted unpickler. Checkpoint format, SHA-256 pins, and all
downstream payload accesses are unchanged.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Anupam Mediratta
2026-07-22 12:28:28 +05:30
co-authored by Claude Sonnet 4.6
parent b2b3851f85
commit 03f1cbb2c2
2 changed files with 9 additions and 7 deletions
+6 -5
View File
@@ -1,8 +1,8 @@
"""Export the pinned JoyVASA diffusion denoiser to a browser-oriented ONNX graph.
The checkpoint pickle global list must be audited before running this script. The
pinned checkpoint currently contains only argparse.Namespace, pathlib.PosixPath,
collections.OrderedDict, and PyTorch tensor rebuild/storage globals.
The checkpoint contains argparse.Namespace, pathlib.PosixPath, collections.OrderedDict,
and PyTorch tensor rebuild/storage globals. The non-default classes are explicitly
allowlisted via add_safe_globals; all others are rejected by weights_only=True.
"""
from __future__ import annotations
@@ -11,7 +11,7 @@ import argparse
import hashlib
import json
import sys
from pathlib import Path
from pathlib import Path, PosixPath
import numpy as np
import onnx
@@ -54,7 +54,8 @@ def main() -> None:
dit_module.enc_dec_mask = cpu_enc_dec_mask
DenoisingNetwork = dit_module.DenoisingNetwork
payload = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
torch.serialization.add_safe_globals([argparse.Namespace, PosixPath])
payload = torch.load(args.checkpoint, map_location="cpu", weights_only=True)
checkpoint_args = payload["args"]
checkpoint_expected = {
"n_diff_steps": 50,
+3 -2
View File
@@ -6,7 +6,7 @@ import argparse
import hashlib
import json
import pickle
from pathlib import Path
from pathlib import Path, PosixPath
import numpy as np
import torch
@@ -32,7 +32,8 @@ def main() -> None:
if sha256(args.template) != EXPECTED_TEMPLATE_SHA256:
raise RuntimeError("Unexpected JoyVASA motion template")
payload = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
torch.serialization.add_safe_globals([argparse.Namespace, PosixPath])
payload = torch.load(args.checkpoint, map_location="cpu", weights_only=True)
# The pinned 3.6KB pickle was audited before use and contains only a dict,
# NumPy ndarray reconstruction, ndarray, and dtype globals.
with args.template.open("rb") as handle: