Validate full safetensors data region is contiguous in no-mmap loader

Per-tensor ranges were checked individually but gaps, overlaps, a
non-zero first offset, and trailing bytes were still accepted. Sort
ranges by start and require the data region to be fully and
contiguously covered, matching safetensors' own validation on the
mmap path. Adds corruption tests for a gap, an overlap, and trailing
bytes.
This commit is contained in:
chelsealong 2026-08-09 01:46:06 +00:00
parent b764c6cbb8
commit 624053ee2f
2 changed files with 65 additions and 2 deletions

View File

@ -135,6 +135,9 @@ def load_safetensors_no_mmap(ckpt, device):
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
data_size = os.fstat(f.fileno()).st_size - data_start
tensors = []
for name, info in header.items():
if name == "__metadata__":
continue
@ -142,9 +145,23 @@ def load_safetensors_no_mmap(ckpt, device):
dtype = _TYPES[info["dtype"]]
shape = info["shape"]
expected_size = math.prod(shape) * dtype.itemsize
if end < start or end - start != expected_size:
if start < 0 or 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:
tensors.append((start, end, name, dtype, shape))
# The data region must be fully and contiguously indexed by the header,
# with no gaps, overlaps, or trailing bytes, matching the validation
# the safetensors library itself performs on the mmap path.
next_start = 0
for start, end, name, _dtype, _shape in sorted(tensors, key=lambda t: t[0]):
if start != next_start:
raise ValueError("Invalid safetensors header: tensor data ranges are not contiguous")
next_start = end
if next_start != data_size:
raise ValueError("Invalid safetensors header: tensor data does not cover the full file")
for start, end, name, dtype, shape in tensors:
if start == end:
sd[name] = torch.empty(shape, dtype=dtype, device=device)
continue
f.seek(data_start + start)

View File

@ -48,3 +48,49 @@ def test_load_safetensors_no_mmap_rejects_corrupt_data_offsets(tmp_path):
with pytest.raises(ValueError):
comfy.utils.load_safetensors_no_mmap(str(path), torch.device("cpu"))
def _write_safetensors_with_raw_offsets(path, header, data):
header_bytes = json.dumps(header).encode("utf-8")
with open(path, "wb") as f:
f.write(struct.pack("<Q", len(header_bytes)))
f.write(header_bytes)
f.write(data)
def test_load_safetensors_no_mmap_rejects_gap_between_tensors(tmp_path):
# weight2 starts one byte after weight ends, leaving an unindexed gap.
header = {
"weight": {"dtype": "F32", "shape": [1], "data_offsets": [0, 4]},
"weight2": {"dtype": "F32", "shape": [1], "data_offsets": [5, 9]},
}
path = tmp_path / "gap.safetensors"
_write_safetensors_with_raw_offsets(path, header, b"\x00" * 9)
with pytest.raises(ValueError):
comfy.utils.load_safetensors_no_mmap(str(path), torch.device("cpu"))
def test_load_safetensors_no_mmap_rejects_overlapping_tensors(tmp_path):
# weight2 starts before weight ends, so the ranges overlap.
header = {
"weight": {"dtype": "F32", "shape": [1], "data_offsets": [0, 4]},
"weight2": {"dtype": "F32", "shape": [1], "data_offsets": [2, 6]},
}
path = tmp_path / "overlap.safetensors"
_write_safetensors_with_raw_offsets(path, header, b"\x00" * 6)
with pytest.raises(ValueError):
comfy.utils.load_safetensors_no_mmap(str(path), torch.device("cpu"))
def test_load_safetensors_no_mmap_rejects_trailing_bytes(tmp_path):
# The declared data region ends before the end of the file.
header = {
"weight": {"dtype": "F32", "shape": [1], "data_offsets": [0, 4]},
}
path = tmp_path / "trailing.safetensors"
_write_safetensors_with_raw_offsets(path, header, b"\x00" * 8)
with pytest.raises(ValueError):
comfy.utils.load_safetensors_no_mmap(str(path), torch.device("cpu"))