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.