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:
co-authored by
Claude Sonnet 4.6
parent
b2b3851f85
commit
03f1cbb2c2
@@ -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,
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user