diff --git a/comfy/utils.py b/comfy/utils.py index 61c2a22dd..0814f7ccd 100644 --- a/comfy/utils.py +++ b/comfy/utils.py @@ -119,6 +119,59 @@ 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 + # violation for large files in a long-running process. Read the tensors + # with plain file I/O instead so no mmap of the file is ever created. + 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 + data_size = os.fstat(f.fileno()).st_size - data_start + + tensors = [] + for name, info in header.items(): + if name == "__metadata__": + continue + start, end = info["data_offsets"] + dtype = _TYPES[info["dtype"]] + shape = info["shape"] + expected_size = math.prod(shape) * dtype.itemsize + 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)) + 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) + 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__", {}) + + def load_torch_file(ckpt, safe_load=False, device=None, return_metadata=False): if device is None: device = torch.device("cpu") @@ -129,14 +182,15 @@ def load_torch_file(ckpt, safe_load=False, device=None, return_metadata=False): sd, metadata = load_safetensors(ckpt) if not return_metadata: metadata = None + elif DISABLE_MMAP: + sd, metadata = load_safetensors_no_mmap(ckpt, device) + if not return_metadata: + metadata = None else: with safetensors.safe_open(ckpt, framework="pt", device=device.type) as f: sd = {} for k in f.keys(): - tensor = f.get_tensor(k) - if DISABLE_MMAP: # TODO: Not sure if this is the best way to bypass the mmap issues - tensor = tensor.to(device=device, copy=True) - sd[k] = tensor + sd[k] = f.get_tensor(k) if return_metadata: metadata = f.metadata() except Exception as e: diff --git a/tests-unit/comfy_test/load_torch_file_test.py b/tests-unit/comfy_test/load_torch_file_test.py new file mode 100644 index 000000000..f64355bd7 --- /dev/null +++ b/tests-unit/comfy_test/load_torch_file_test.py @@ -0,0 +1,96 @@ +import json +import os +import struct +import tempfile + +import pytest +import safetensors.torch +import torch + +import comfy.utils + + +@pytest.fixture +def safetensors_file(): + tensors = {"weight": torch.arange(12, dtype=torch.float32).reshape(3, 4)} + with tempfile.TemporaryDirectory() as tmpdirname: + path = os.path.join(tmpdirname, "model.safetensors") + safetensors.torch.save_file(tensors, path, metadata={"format": "pt"}) + yield path, tensors + + +def test_disable_mmap_does_not_use_safe_open(safetensors_file, monkeypatch): + path, tensors = safetensors_file + monkeypatch.setattr(comfy.utils, "DISABLE_MMAP", True) + + def boom(*args, **kwargs): + raise AssertionError("safetensors.safe_open should not be used when DISABLE_MMAP is set") + + monkeypatch.setattr(comfy.utils.safetensors, "safe_open", boom) + + sd, metadata = comfy.utils.load_torch_file(path, device=torch.device("cpu"), return_metadata=True) + + 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("