diodiogod
58678e2497
- Move Shared Runtime preparation to the installation hook and add persistent UI/file opt-out. - Restrict file and model download paths, checkpoint loading, and configuration parsing across the suite. - Remove Dedicated Runtime choices and migrate saved workflows without shifting widgets. - Exclude development files from Registry archives while retaining runtime metadata and Unicode paths. - Document installation, repair, and manual runtime cleanup. Validation: - 59 targeted existing tests, actual fresh runtime installation, and live Step speech/editing. - File containment, restricted loader compatibility, settings persistence, and archive contents checked. - Manager queue/restart flow remains untested on the local installation. Address Registry availability reported in #366.
88 lines
3.2 KiB
Python
88 lines
3.2 KiB
Python
"""Load tensor checkpoints and plain dictionaries without executable pickle globals."""
|
|
|
|
import pickle
|
|
|
|
|
|
class DataOnlyUnpickler(pickle.Unpickler):
|
|
def find_class(self, module, name):
|
|
raise pickle.UnpicklingError(f"Executable pickle global is forbidden: {module}.{name}")
|
|
|
|
|
|
def load_data_pickle(handle):
|
|
return DataOnlyUnpickler(handle).load()
|
|
|
|
|
|
class NumpyDataUnpickler(DataOnlyUnpickler):
|
|
def find_class(self, module, name):
|
|
import numpy as np
|
|
|
|
core = np._core if hasattr(np, "_core") else np.core
|
|
allowed = {
|
|
("numpy", "ndarray"): np.ndarray,
|
|
("numpy", "dtype"): np.dtype,
|
|
("numpy.core.multiarray", "_reconstruct"): core.multiarray._reconstruct,
|
|
("numpy._core.multiarray", "_reconstruct"): core.multiarray._reconstruct,
|
|
("numpy.core.multiarray", "scalar"): core.multiarray.scalar,
|
|
("numpy._core.multiarray", "scalar"): core.multiarray.scalar,
|
|
}
|
|
if (module, name) in allowed:
|
|
return allowed[module, name]
|
|
return super().find_class(module, name)
|
|
|
|
|
|
def load_numpy_data(path):
|
|
"""Keep .npy/.dac data dictionaries while refusing arbitrary pickle globals."""
|
|
import numpy as np
|
|
|
|
with open(path, "rb") as handle:
|
|
version = np.lib.format.read_magic(handle)
|
|
shape, _fortran, dtype = np.lib.format._read_array_header(handle, version)
|
|
if not dtype.hasobject:
|
|
handle.seek(0)
|
|
return np.load(handle, allow_pickle=False)
|
|
result = NumpyDataUnpickler(handle).load()
|
|
if not isinstance(result, np.ndarray) or result.shape != shape or result.dtype != dtype:
|
|
raise ValueError("NumPy data does not match its array header.")
|
|
return result
|
|
|
|
|
|
def load_torch_checkpoint(*args, **kwargs):
|
|
import numpy as np
|
|
import torch
|
|
|
|
if kwargs.get("weights_only") is False or "pickle_module" in kwargs:
|
|
raise ValueError("Checkpoint loading must use weights_only=True.")
|
|
kwargs["weights_only"] = True
|
|
# Dots latent statistics contain NumPy arrays. This fixed, scoped allowlist
|
|
# supports those data values without accepting arbitrary checkpoint classes.
|
|
numpy_core = np._core if hasattr(np, "_core") else np.core
|
|
safe_types = [
|
|
numpy_core.multiarray._reconstruct,
|
|
(numpy_core.multiarray._reconstruct, "numpy.core.multiarray._reconstruct"),
|
|
np.ndarray,
|
|
np.dtype,
|
|
type(np.dtype("float32")),
|
|
type(np.dtype("float64")),
|
|
]
|
|
with torch.serialization.safe_globals(safe_types):
|
|
return torch.load(*args, **kwargs)
|
|
|
|
|
|
class RestrictedLoadModule:
|
|
"""Restrict a dependency's own module reference without patching Python globally."""
|
|
|
|
def __init__(self, module, load=None, **overrides):
|
|
self._module = module
|
|
self._overrides = overrides
|
|
if load is not None:
|
|
self._overrides["load"] = load
|
|
|
|
def __getattr__(self, name):
|
|
return self._overrides[name] if name in self._overrides else getattr(self._module, name)
|
|
|
|
|
|
def load_dependency_checkpoint(*args, **kwargs):
|
|
# Override a dependency's explicit unrestricted flag in this integration.
|
|
kwargs["weights_only"] = True
|
|
return load_torch_checkpoint(*args, **kwargs)
|