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:
parent
b764c6cbb8
commit
624053ee2f
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"))
|
||||
|
|
|
|||
Loading…
Reference in New Issue