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

@@ -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("<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
         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__", {})
 
 

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("<Q", len(header_bytes)))
+        f.write(header_bytes)
+
+    with pytest.raises(ValueError):
+        comfy.utils.load_safetensors_no_mmap(str(path), torch.device("cpu"))