From 5dfb3002ea64f9be0bb66d3df1fe1849a6f5ba57 Mon Sep 17 00:00:00 2001 From: chelsealong Date: Sun, 9 Aug 2026 01:18:48 +0000 Subject: [PATCH 1/3] Fix --disable-mmap not bypassing safetensors mmap read load_torch_file only wrapped the tensor returned by safetensors.safe_open()/get_tensor() in a copy when --disable-mmap was set; the mmap-backed read that produces that tensor had already happened by then. On Windows that mmap-backed read of large safetensors files can crash the long-running ComfyUI process with an access violation (#15424). --disable-mmap now reads tensors with plain file I/O so no mmap of the file is created at all. --- comfy/utils.py | 34 ++++++++++++++++--- tests-unit/comfy_test/load_torch_file_test.py | 32 +++++++++++++++++ 2 files changed, 62 insertions(+), 4 deletions(-) create mode 100644 tests-unit/comfy_test/load_torch_file_test.py diff --git a/comfy/utils.py b/comfy/utils.py index 61c2a22dd..b159844c9 100644 --- a/comfy/utils.py +++ b/comfy/utils.py @@ -119,6 +119,31 @@ def load_safetensors(ckpt): return sd, header.get("__metadata__", {}), +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(" Date: Sun, 9 Aug 2026 01:35:01 +0000 Subject: [PATCH 2/3] Validate safetensors header/offsets and avoid double buffer in no-mmap loader Address CodeRabbit review on the --disable-mmap fix: reject headers over the safetensors size limit and tensor data ranges that don't match their declared shape/dtype (previously a corrupt header with start==end for a non-empty shape produced an uninitialized tensor instead of an error), and read each tensor directly into a single bytearray via readinto() instead of read() + bytearray() copy to halve peak host memory use for large files. --- comfy/utils.py | 17 ++++++++++++++--- tests-unit/comfy_test/load_torch_file_test.py | 18 ++++++++++++++++++ 2 files changed, 32 insertions(+), 3 deletions(-) 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(" Date: Sun, 9 Aug 2026 01:46:06 +0000 Subject: [PATCH 3/3] 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. --- comfy/utils.py | 21 ++++++++- tests-unit/comfy_test/load_torch_file_test.py | 46 +++++++++++++++++++ 2 files changed, 65 insertions(+), 2 deletions(-) diff --git a/comfy/utils.py b/comfy/utils.py index 175838547..0814f7ccd 100644 --- a/comfy/utils.py +++ b/comfy/utils.py @@ -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) diff --git a/tests-unit/comfy_test/load_torch_file_test.py b/tests-unit/comfy_test/load_torch_file_test.py index eddd12a87..f64355bd7 100644 --- a/tests-unit/comfy_test/load_torch_file_test.py +++ b/tests-unit/comfy_test/load_torch_file_test.py @@ -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("