Merge 624053ee2f into 37ac9ff44f

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("<Q", f.read(8))[0]
+        if header_size > 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:

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("<Q", len(header_bytes)))
+        f.write(header_bytes)
+
+    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"))