From 57f816b4d4b71371739155a1b85f56f653cdfc35 Mon Sep 17 00:00:00 2001 From: Asish Kumar Date: Mon, 25 May 2026 05:48:01 +0530 Subject: [PATCH] fix: load voice presets with a restricted unpickler (CWE-502) Commit 303b283 switched the voice-preset loaders in demo/web/app.py and demo/realtime_model_inference_from_file.py to `torch.load(..., weights_only=True)` guarded by `torch.serialization.safe_globals([BaseModelOutputWithPast, DynamicCache])` to close a CWE-502 arbitrary-code-execution risk. That load can never succeed. The presets are dicts of `transformers.modeling_outputs.BaseModelOutputWithPast` objects, and PyTorch's weights-only unpickler refuses the `SETITEMS` opcode on `dict` subclasses (it accepts only the exact `dict`, `OrderedDict` and `Counter` types) even when the class is allowlisted via `safe_globals`. Loading therefore aborts on startup with: _pickle.UnpicklingError: Weights only load failed. ... Can only SETITEMS for dict, collections.OrderedDict, collections.Counter, but got making the streaming web demo and the realtime file demo unusable on PyTorch >= 2.6. `BaseModelOutputWithPast` has to stay an object (the model accesses `outputs.past_key_values` and `outputs.last_hidden_state`), so the presets cannot simply be flattened to plain tensors. Instead, load them through a restricted `pickle.Unpickler` whose `find_class` resolves only the container classes the presets are built from (`OrderedDict`, `BaseModelOutputWithPast`, `DynamicCache`) plus torch's tensor-rebuilding primitives. Any other global, e.g. `os.system`, is refused, so a tampered preset still cannot execute arbitrary code and the CWE-502 protection is preserved. The same restriction is applied to the module-level `load` / `loads` that `torch.load` uses for legacy-format metadata, so both serialization formats stay safe. The shared `load_voice_preset` helper lives in the streaming processor module imported by both demos; the now-unused `BaseModelOutputWithPast` and `DynamicCache` imports are removed. --- demo/realtime_model_inference_from_file.py | 7 +- demo/web/app.py | 12 ++-- .../vibevoice_streaming_processor.py | 67 +++++++++++++++++++ 3 files changed, 73 insertions(+), 13 deletions(-) diff --git a/demo/realtime_model_inference_from_file.py b/demo/realtime_model_inference_from_file.py index 2a2e711c..714771e7 100644 --- a/demo/realtime_model_inference_from_file.py +++ b/demo/realtime_model_inference_from_file.py @@ -9,9 +9,7 @@ import glob from vibevoice.modular.modeling_vibevoice_streaming_inference import VibeVoiceStreamingForConditionalGenerationInference -from vibevoice.processor.vibevoice_streaming_processor import VibeVoiceStreamingProcessor -from transformers.cache_utils import DynamicCache -from transformers.modeling_outputs import BaseModelOutputWithPast +from vibevoice.processor.vibevoice_streaming_processor import VibeVoiceStreamingProcessor, load_voice_preset from transformers.utils import logging logging.set_verbosity_info() @@ -224,8 +222,7 @@ def main(): target_device = args.device if args.device != "cpu" else "cpu" voice_sample = voice_mapper.get_voice_path(args.speaker_name) print(f"Using voice preset for {args.speaker_name}: {voice_sample}") - with torch.serialization.safe_globals([BaseModelOutputWithPast, DynamicCache]): - all_prefilled_outputs = torch.load(voice_sample, map_location=target_device, weights_only=True) + all_prefilled_outputs = load_voice_preset(voice_sample, map_location=target_device) # Prepare inputs for the model inputs = processor.process_input_with_cached_prompt( diff --git a/demo/web/app.py b/demo/web/app.py index 4985c7db..86aac381 100644 --- a/demo/web/app.py +++ b/demo/web/app.py @@ -11,8 +11,6 @@ import numpy as np import torch -from transformers.cache_utils import DynamicCache -from transformers.modeling_outputs import BaseModelOutputWithPast from fastapi import FastAPI, WebSocket from fastapi.responses import FileResponse from fastapi.staticfiles import StaticFiles @@ -23,6 +21,7 @@ ) from vibevoice.processor.vibevoice_streaming_processor import ( VibeVoiceStreamingProcessor, + load_voice_preset, ) from vibevoice.modular.streamer import AudioStreamer @@ -160,12 +159,9 @@ def _ensure_voice_cached(self, key: str) -> Tuple[object, Path, str]: preset_path = self.voice_presets[key] print(f"[startup] Loading voice preset {key} from {preset_path}") print(f"[startup] Loading prefilled prompt from {preset_path}") - with torch.serialization.safe_globals([BaseModelOutputWithPast, DynamicCache]): - prefilled_outputs = torch.load( - preset_path, - map_location=self._torch_device, - weights_only=True, - ) + prefilled_outputs = load_voice_preset( + preset_path, map_location=self._torch_device + ) self._voice_cache[key] = prefilled_outputs return self._voice_cache[key] diff --git a/vibevoice/processor/vibevoice_streaming_processor.py b/vibevoice/processor/vibevoice_streaming_processor.py index 39c262b1..d4b69bd3 100644 --- a/vibevoice/processor/vibevoice_streaming_processor.py +++ b/vibevoice/processor/vibevoice_streaming_processor.py @@ -1,4 +1,6 @@ +import io import math +import pickle import warnings from typing import List, Optional, Union, Dict, Any, Tuple import os @@ -14,6 +16,71 @@ logger = logging.get_logger(__name__) +# Globals that legitimately appear in the cached voice-prompt (.pt) files. Any +# other global (e.g. os.system) is refused, so a tampered preset cannot execute +# arbitrary code when it is loaded (CWE-502). +_VOICE_PRESET_SAFE_GLOBALS = { + ("collections", "OrderedDict"), + ("transformers.modeling_outputs", "BaseModelOutputWithPast"), + ("transformers.cache_utils", "DynamicCache"), +} + + +class _VoicePresetUnpickler(pickle.Unpickler): + """Unpickler that only resolves the classes used by voice presets.""" + + def find_class(self, module, name): + if (module, name) in _VOICE_PRESET_SAFE_GLOBALS: + return super().find_class(module, name) + # torch tensor-rebuilding primitives; the storages they reference are + # restored by torch.load's persistent_load, not through find_class. + if module == "torch._utils" and name.startswith("_rebuild_"): + return super().find_class(module, name) + if module == "torch" and name.endswith("Storage"): + return super().find_class(module, name) + raise pickle.UnpicklingError( + f"Refusing to load disallowed global '{module}.{name}' from voice preset" + ) + + +class _RestrictedPickleModule: + """``pickle_module`` exposing the restricted unpickler to ``torch.load``. + + ``torch.load`` uses ``Unpickler`` for the main payload and the module-level + ``load``/``loads`` for legacy-format metadata, so all three must enforce the + same restriction. + """ + + Unpickler = _VoicePresetUnpickler + + @staticmethod + def load(file, **kwargs): + return _VoicePresetUnpickler(file, **kwargs).load() + + @staticmethod + def loads(data, **kwargs): + return _VoicePresetUnpickler(io.BytesIO(data), **kwargs).load() + + +def load_voice_preset(path, map_location=None): + """Load a cached voice-prompt file produced for streaming TTS. + + The presets store ``transformers`` ``BaseModelOutputWithPast`` / + ``DynamicCache`` objects (``dict`` subclasses), which + ``torch.load(weights_only=True)`` cannot rebuild because its unpickler + forbids ``SETITEMS`` on ``dict`` subclasses. A restricted unpickler is used + instead: it permits only those container classes and torch's tensor + primitives, keeping the protection against arbitrary code execution that + ``weights_only=True`` was added for (CWE-502). + """ + return torch.load( + path, + map_location=map_location, + pickle_module=_RestrictedPickleModule, + weights_only=False, + ) + + class VibeVoiceStreamingProcessor: r""" Constructs a VibeVoice Streaming processor which wraps a VibeVoice tokenizer and audio processor into a single processor.