fix: load voice presets with a restricted unpickler (CWE-502)

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(

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]

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.