diff --git a/comfy/utils.py b/comfy/utils.py index b159844c9..175838547 100644 --- a/comfy/utils.py +++ b/comfy/utils.py @@ -119,6 +119,10 @@ def load_safetensors(ckpt): return sd, header.get("__metadata__", {}), +# Matches the header size limit enforced by the safetensors library itself. +MAX_SAFETENSORS_HEADER_SIZE = 100_000_000 + + def load_safetensors_no_mmap(ckpt, device): # safetensors.safe_open()/get_tensor() reads tensor data through an mmap of # the file. On Windows that mmap-backed read can crash with an access @@ -127,6 +131,8 @@ def load_safetensors_no_mmap(ckpt, device): sd = {} with open(ckpt, "rb") as f: header_size = struct.unpack(" MAX_SAFETENSORS_HEADER_SIZE: + raise ValueError("Invalid safetensors header: header size exceeds the maximum allowed size") header = json.loads(f.read(header_size).decode("utf-8")) data_start = 8 + header_size for name, info in header.items(): @@ -135,12 +141,17 @@ def load_safetensors_no_mmap(ckpt, device): start, end = info["data_offsets"] dtype = _TYPES[info["dtype"]] shape = info["shape"] - if start == end: + expected_size = math.prod(shape) * dtype.itemsize + if end < start or end - start != expected_size: + raise ValueError("Invalid safetensors header: tensor '{}' data range does not match its declared shape/dtype".format(name)) + if expected_size == 0: sd[name] = torch.empty(shape, dtype=dtype, device=device) continue f.seek(data_start + start) - raw = f.read(end - start) - sd[name] = torch.frombuffer(bytearray(raw), dtype=dtype).view(shape).to(device=device) + raw = bytearray(end - start) + if f.readinto(raw) != len(raw): + raise ValueError("Invalid safetensors file: tensor '{}' data is truncated".format(name)) + sd[name] = torch.frombuffer(raw, dtype=dtype).view(shape).to(device=device) return sd, header.get("__metadata__", {}) diff --git a/tests-unit/comfy_test/load_torch_file_test.py b/tests-unit/comfy_test/load_torch_file_test.py index b220f8f6a..eddd12a87 100644 --- a/tests-unit/comfy_test/load_torch_file_test.py +++ b/tests-unit/comfy_test/load_torch_file_test.py @@ -1,4 +1,6 @@ +import json import os +import struct import tempfile import pytest @@ -30,3 +32,19 @@ def test_disable_mmap_does_not_use_safe_open(safetensors_file, monkeypatch): assert torch.equal(sd["weight"], tensors["weight"]) assert metadata == {"format": "pt"} + + +def test_load_safetensors_no_mmap_rejects_corrupt_data_offsets(tmp_path): + # A corrupt header claiming a non-empty shape with start == end must be + # rejected instead of silently producing an uninitialized tensor. + header = { + "weight": {"dtype": "F32", "shape": [3, 4], "data_offsets": [0, 0]}, + } + header_bytes = json.dumps(header).encode("utf-8") + path = tmp_path / "corrupt.safetensors" + with open(path, "wb") as f: + f.write(struct.pack("