51 lines
1.6 KiB
Python
51 lines
1.6 KiB
Python
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"))
|