diff --git a/comfy/utils.py b/comfy/utils.py index 61c2a22dd..72327b053 100644 --- a/comfy/utils.py +++ b/comfy/utils.py @@ -96,6 +96,11 @@ def load_safetensors(ckpt): mv = mv[(data_base_offset := 8 + header_size):] + expected_data_size = max((info["data_offsets"][1] for name, info in header.items() if name != "__metadata__"), default=0) + if len(mv) < expected_data_size: + message = "buffer length ({} bytes) is smaller than the {} bytes of tensor data declared in the header".format(len(mv), expected_data_size) + raise ValueError("{}\n\nFile path: {}\n\nThe safetensors file is corrupt/incomplete. Check the file size and make sure you have copied/downloaded it correctly.".format(message, ckpt)) + sd = {} for name, info in header.items(): if name == "__metadata__": diff --git a/tests-unit/comfy_test/load_safetensors_test.py b/tests-unit/comfy_test/load_safetensors_test.py new file mode 100644 index 000000000..e4e88126c --- /dev/null +++ b/tests-unit/comfy_test/load_safetensors_test.py @@ -0,0 +1,115 @@ +""" +Regression test for https://github.com/comfyanonymous/ComfyUI/issues/14784 + +VAELoader (and any other model loader going through comfy.utils.load_torch_file) +raised a low-level, confusing: + + ValueError: buffer length (N bytes) after offset (0 bytes) must be a + multiple of element size (4) + +instead of a clear "corrupt/incomplete file" message, whenever the aimdo +mmap-based safetensors loader (comfy.utils.load_safetensors) was handed a +truncated/incomplete .safetensors file (e.g. one that was still downloading, +or whose download was interrupted). The root cause: load_safetensors() never +validated that the mapped file actually contained as many bytes as the +header declared before slicing them and handing the (silently truncated) +slice to torch.frombuffer(). + +comfy_aimdo.model_mmap.ModelMMAP requires a compiled, GPU-platform-specific +native backend that is unavailable on this (and any non-GPU CI) machine, so +it is monkeypatched here with a fake that maps a real, small, synthetic +safetensors-formatted buffer we fully control. This lets the test exercise +comfy.utils.load_safetensors() itself, deterministically and without any +real checkpoint or GPU. +""" +import ctypes +import json +import os +import struct +import sys + +import pytest + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..")) + +from comfy.cli_args import args # noqa: E402 + +if not __import__("torch").cuda.is_available(): + args.cpu = True + +import comfy.utils # noqa: E402 +import comfy_aimdo.model_mmap # noqa: E402 + + +def _write_fake_safetensors(path, data_size, declared_size): + """Write a minimal safetensors-formatted file with a single F32 tensor + whose header declares `declared_size` bytes of data, but only + `data_size` bytes are actually written after the header (simulating a + truncated/incomplete download when data_size < declared_size).""" + header = {"test.weight": {"dtype": "F32", "shape": [declared_size // 4], "data_offsets": [0, declared_size]}} + header_bytes = json.dumps(header).encode("utf-8") + with open(path, "wb") as f: + f.write(struct.pack("