invoke-ai/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI.git synced 2026-10-12 02:16:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

test(fixtures): capture a key layout from the safetensors header alone

Browse Source 
Replaces scripts/extract_sd_keys_and_shapes.py, which read whole checkpoints
through load_file and recorded shapes only. Reads the header through a local
handle or an HTTP range request, one parser over both, and records dtypes, the
_quantization_metadata block and the decoded .comfy_quant markers.
This commit is contained in:
Alexander Eichhorn
2026-09-26 04:40:23 +02:00
parent 4401030b25
commit 2edfb3e2df
4 changed files with 781 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/capture_state_dict_fixture.py
+413
View File
@@ -0,0 +1,413 @@
"""Capture a checkpoint's key layout as a test fixture module, without reading its weights.
The loader conformance work needs real key layouts: invented names fall through the converters'
format probes as "no match", the converter becomes a no-op, and a test written against them asserts
nothing. But everything such a fixture records -- key names, shapes, dtypes, scale spellings, the
`_quantization_metadata` header block and the per-layer `.comfy_quant` markers -- lives in the
safetensors *header*, which is a few kilobytes at the front of the file. So a layout nobody here
owns is still a layout that can be reproduced: point this at a URL and it reads the header over an
HTTP range request.
Two byte sources, one parser. The local source is deliberately not `safe_open`: a second reading
path would leave the remote one -- the only one that can reach a file we do not have -- as the
untested half.
The output is a Python module in the shape of the fixtures beside
`tests/backend/model_manager/load/state_dicts/`, not JSON, because that is what the tests import.
Its docstring is a stub: it records where the layout came from and what was kept, and leaves the
part no tool can write -- what the layout *means*, and which loader property depends on it. Run
`ruff format` on the result; the checked-in fixtures are formatted.
Usage:
uv run --no-sync python scripts/capture_state_dict_fixture.py <source> --out <module.py>
`<source>` is a local path, a full `https://...` URL, or `<org>/<repo>/<path-in-repo>`, which is
resolved against Hugging Face. A local path wins when the file exists.
"""
from __future__ import annotations
import argparse
import datetime
import json
import re
import struct
import sys
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable, Mapping, Protocol
import httpx
from huggingface_hub.utils import build_hf_headers
from invokeai.backend.quantization.fp8_scaled import COMFY_QUANT_SUFFIX, parse_quantization_metadata
from invokeai.backend.quantization.int8_convrot import parse_comfy_quant_bytes
# Not read from the environment on purpose: `HF_ENDPOINT` often points at a caching proxy here, and
# proxies answer range reads inconsistently (401 on some Comfy-Org files). `build_hf_headers()` does
# read the environment, which is what supplies the token.
HF_ENDPOINT = "https://huggingface.co"
# A sanity bound on the declared header length. The largest checkpoint reachable from this tree
# declares well under a megabyte; anything past this is a file that is not a safetensors.
MAX_HEADER_BYTES = 64 * 1024 * 1024
class ByteSource(Protocol):
"""A file the header can be read out of, byte range by byte range."""
def read(self, offset: int, length: int) -> bytes: ...
def close(self) -> None: ...
class LocalFile:
"""A file on disk. The handle stays open: a capture with markers does one read per marker."""
def __init__(self, path: Path) -> None:
self._handle = path.open("rb")
self.display = path.as_posix()
def read(self, offset: int, length: int) -> bytes:
self._handle.seek(offset)
return self._handle.read(length)
def close(self) -> None:
self._handle.close()
def __str__(self) -> str:
return self.display
class RemoteFile:
"""A file behind an HTTP server that honours `Range`, which Hugging Face's CDN does.
Nothing but the requested ranges is transferred, so a 20 GiB checkpoint costs a few kilobytes.
"""
def __init__(self, url: str, display: str | None = None) -> None:
self._url = url
self.display = display or url
self._client = httpx.Client(follow_redirects=True, timeout=180, headers=build_hf_headers())
def read(self, offset: int, length: int) -> bytes:
response = self._client.get(self._url, headers={"Range": f"bytes={offset}-{offset + length - 1}"})
response.raise_for_status()
if response.status_code != httpx.codes.PARTIAL_CONTENT:
# 200 means the server ignored `Range` and sent the body -- for a 14 GB checkpoint that
# is the entire download this tool exists to avoid, already buffered in memory by the
# time we could notice. Refuse rather than continue: the capture would "work".
raise SystemExit(
f"{self._url} answered {response.status_code} to a range request instead of 206; "
"the server or a redirect target does not honour `Range`, so reading its header "
"would mean downloading the whole file"
)
return response.content
def close(self) -> None:
self._client.close()
def __str__(self) -> str:
return self.display
def _looks_like_a_path(source: str) -> bool:
"""Whether a source that does not exist was nonetheless meant as a file rather than a repo.
Without this a typo in a local path has three or more segments and falls through to the
repository branch, so `S:/ModelStuff3/typo.safetensors` is fetched from
`huggingface.co/S:/ModelStuff3/...` and the operator is told about an HTTP status instead of a
missing file. A repository id has none of these shapes.
"""
return source.startswith(("/", "./", "../", "~")) or "\\" in source or (len(source) > 1 and source[1] == ":")
def open_source(source: str) -> ByteSource:
"""A local path if it exists, a URL as given, otherwise `<org>/<repo>/<path>` on Hugging Face."""
if source.startswith(("http://", "https://")):
return RemoteFile(source)
path = Path(source)
if path.exists():
return LocalFile(path)
if _looks_like_a_path(source):
raise SystemExit(f"no such file: {source}")
parts = source.split("/")
if len(parts) < 3:
raise SystemExit(f"{source!r} is neither an existing path, a URL, nor <org>/<repo>/<path-in-repo>")
repo, in_repo = "/".join(parts[:2]), "/".join(parts[2:])
# The fixture cites the repository, not the resolve URL: that is how a reader finds the file
# again, and it is the form the hand-written fixtures beside it already use.
return RemoteFile(f"{HF_ENDPOINT}/{repo}/resolve/main/{in_repo}", display=source)
@dataclass(frozen=True)
class TensorEntry:
"""One header entry. `begin`/`end` are absolute file offsets, not the header's relative ones."""
dtype: str
shape: list[int]
begin: int
end: int
def read_header(source: ByteSource) -> tuple[dict[str, TensorEntry], dict[str, str]]:
"""The safetensors header: eight bytes of little-endian length, then that much JSON."""
(header_length,) = struct.unpack("<Q", source.read(0, 8))
if not 0 < header_length <= MAX_HEADER_BYTES:
# The length is read from the file, so a URL that is not a safetensors -- an HTML error
# page, an LFS pointer -- gives a nonsense one, and the next read asks for that many bytes.
raise SystemExit(
f"{source} declares a {header_length}-byte header; that is not a safetensors file "
f"(the limit here is {MAX_HEADER_BYTES} bytes, and the largest checkpoint in this tree "
"declares under a megabyte)"
)
header = json.loads(source.read(8, header_length))
metadata = header.pop("__metadata__", None) or {}
data_start = 8 + header_length
tensors = {
key: TensorEntry(
dtype=entry["dtype"],
shape=list(entry["shape"]),
begin=data_start + entry["data_offsets"][0],
end=data_start + entry["data_offsets"][1],
)
for key, entry in header.items()
}
return tensors, metadata
def read_markers(source: ByteSource, tensors: Mapping[str, TensorEntry]) -> dict[str, dict[str, Any]]:
"""Decode every `<layer>.comfy_quant` blob, keyed by the layer path it names.
These are 64 to 89 bytes each, so reading them is cheap even over the network -- and it is the
difference between a fixture that records the marker and one that only records its length.
"""
markers: dict[str, dict[str, Any]] = {}
for key, entry in tensors.items():
if not key.endswith(COMFY_QUANT_SUFFIX):
continue
markers[key[: -len(COMFY_QUANT_SUFFIX)]] = parse_comfy_quant_bytes(
source.read(entry.begin, entry.end - entry.begin)
)
return markers
_INDEX = re.compile(r"^\d+$")
# How many distinct indices a prefix needs before it counts as a repeated block stack rather than
# an `nn.Sequential`. Four separates every stack in the captured checkpoints from every Sequential.
MIN_STACK = 4
def stack_indices(keys: Iterable[str]) -> dict[str, set[int]]:
"""Every integer path segment, grouped by the path that precedes it.
A checkpoint has more than one stack -- FLUX keeps `double_blocks` and `single_blocks` -- and
they are numbered independently, so the grouping is by prefix rather than by depth.
"""
found: dict[str, set[int]] = defaultdict(set)
for key in keys:
parts = key.split(".")
for position, part in enumerate(parts):
if _INDEX.match(part):
found[".".join(parts[:position])].add(int(part))
return found
def choose_indices(keys: Iterable[str], block: int | None) -> dict[str, int]:
"""Which index to keep per stack.
Not every integer segment is a stack. `nn.Sequential` numbers its children too, and those are
parts of one module rather than repetitions of it -- Krea-2's `tmlp.0`/`tmlp.2` and Anima's
`adaln_modulation_cross_attn.1`/`.2` are both halves of a thing, and a capture that kept only
the first would describe a module the file does not contain. A repeated block stack is long, a
Sequential is two or three deep, so `MIN_STACK` separates them. Where that heuristic is wrong --
a decoder with four or more Sequential children -- `--no-subset` keeps the file whole, and the
rendered docstring always names which prefixes were treated as stacks, so the call is visible.
`block` is a preference, not a demand: a stack that does not have it keeps its lowest index and
says so. Stacks are different lengths -- asking for block 24 of a 38-deep stack is reasonable
and of a 19-deep one is not -- and refusing the whole capture over that would help nobody.
"""
chosen: dict[str, int] = {}
for prefix, indices in stack_indices(keys).items():
if len(indices) < MIN_STACK:
continue
if block is not None and block not in indices:
print(
f"note: {prefix or '<root>'} has no index {block}; keeping {min(indices)} instead",
file=sys.stderr,
)
chosen[prefix] = block if block is not None and block in indices else min(indices)
return chosen
def select_keys(keys: Iterable[str], chosen: Mapping[str, int]) -> list[str]:
"""Keep the chosen index of every stack, and every key that is not in one."""
kept = []
for key in keys:
parts = key.split(".")
if all(
int(part) == chosen[".".join(parts[:position])]
for position, part in enumerate(parts)
if _INDEX.match(part) and ".".join(parts[:position]) in chosen
):
kept.append(key)
return kept
def _py(value: Any) -> str:
"""Render a JSON-ish value as a Python literal. `bool` before `int`: it is a subclass."""
if isinstance(value, str):
return json.dumps(value)
if isinstance(value, bool):
return "True" if value else "False"
if value is None:
return "None"
if isinstance(value, (int, float)):
# `json.loads` accepts NaN/Infinity, whose `repr` is a bare `nan`/`inf` -- a NameError when
# the fixture is imported. No producer writes them; refusing beats emitting a broken module.
if isinstance(value, float) and (value != value or value in (float("inf"), float("-inf"))):
raise SystemExit(f"header carries a non-finite value ({value!r}), which no fixture can record")
return repr(value)
if isinstance(value, (list, tuple)):
return "[" + ", ".join(_py(item) for item in value) + "]"
if isinstance(value, dict):
return "{" + ", ".join(f"{_py(key)}: {_py(item)}" for key, item in value.items()) + "}"
raise TypeError(f"cannot render {type(value).__name__} into a fixture module")
def _docstring_safe(text: str) -> str:
"""Neutralise anything in a source name that a docstring would read as syntax."""
return text.replace("\\", "\\\\").replace('"""', '\\"\\"\\"')
def render_module(
*,
source: str,
tensors: Mapping[str, TensorEntry],
chosen: Mapping[str, int],
layer_hints: Mapping[str, Any],
markers: Mapping[str, Any],
) -> str:
"""The fixture module, in the shape the state_dicts fixtures beside it already have."""
if chosen:
rule = ", ".join(f"`{prefix or '<root>'}` index {index}" for prefix, index in sorted(chosen.items()))
rule = f"{rule}, plus every key outside a stack"
else:
rule = "none -- every key in the file"
lines = [
'"""TODO: one line naming the build this is the key layout of.',
"",
f"Captured on {datetime.date.today().isoformat()} with `scripts/capture_state_dict_fixture.py`,",
"which reads the header only, from",
f"`{_docstring_safe(source)}`.",
"",
f"Subsetting rule: {rule}. Values are `(shape, dtype)`.",
"",
"TODO: what this layout *means* -- which loader property depends on it, and what a re-upload",
"could change without any loader code moving. A fixture without that is a wall of keys.",
'"""',
"",
"state_dict_keys: dict[str, tuple[list[int], str]] = {",
]
for key in sorted(tensors):
entry = tensors[key]
lines.append(f" {_py(key)}: ({_py(entry.shape)}, {_py(entry.dtype)}),")
lines.append("}")
if layer_hints:
lines += [
"",
"# The `_quantization_metadata` header block, layer names exactly as the producer wrote them.",
"layer_hints: dict[str, dict[str, object]] = {",
]
lines += [f" {_py(name)}: {_py(hint)}," for name, hint in sorted(layer_hints.items())]
lines.append("}")
if markers:
lines += [
"",
"# The per-layer `.comfy_quant` markers, decoded from the blobs themselves.",
"markers: dict[str, dict[str, object]] = {",
]
lines += [f" {_py(name)}: {_py(marker)}," for name, marker in sorted(markers.items())]
lines.append("}")
return "\n".join(lines) + "\n"
def capture(source_argument: str, *, block: int | None, subset: bool, want_markers: bool) -> str:
source = open_source(source_argument)
try:
tensors, metadata = read_header(source)
chosen = choose_indices(tensors, block) if subset else {}
kept = set(select_keys(tensors, chosen)) if chosen else set(tensors)
subsetted = {key: entry for key, entry in tensors.items() if key in kept}
markers = read_markers(source, subsetted) if want_markers else {}
finally:
source.close()
hints = parse_quantization_metadata(metadata)
layer_hints = {name: hint for name, hint in hints.items() if _hint_survived(name, kept)}
if hints and not any(_hint_survived(name, set(tensors)) for name in hints):
# Not the subset's doing: no header name matches any tensor in the whole file. That is the
# header-to-tensor namespace drift `strip_layer_path_prefix` exists for, and a capture that
# said nothing would record the build as having no header block at all.
print(
f"note: none of the {len(hints)} header layer names match a tensor in this file; "
"the fixture will record no layer_hints",
file=sys.stderr,
)
return render_module(
source=str(source),
tensors=subsetted,
chosen=chosen,
layer_hints=layer_hints,
markers=markers,
)
def _hint_survived(name: str, kept: set[str]) -> bool:
"""A header hint names a module, not a tensor, so it is kept when any of its tensors was.
The header keeps the producer's own naming, which is the whole point of recording it: it is the
namespace the loader has to re-key from, and a capture that silently dropped the mismatching
ones would hide exactly the drift the fixture exists to show.
"""
return any(key == name or key.startswith(f"{name}.") for key in kept)
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("source", help="local path, https:// URL, or <org>/<repo>/<path-in-repo>")
parser.add_argument("--out", type=Path, help="module to write; prints to stdout when omitted")
parser.add_argument(
"--block",
type=int,
default=None,
help="stack index to keep (default: the lowest each stack has)",
)
parser.add_argument("--no-subset", action="store_true", help="keep every key, stacks and all")
parser.add_argument("--no-markers", action="store_true", help="skip the .comfy_quant blob reads")
args = parser.parse_args()
module = capture(
args.source,
block=args.block,
subset=not args.no_subset,
want_markers=not args.no_markers,
)
if args.out is None:
print(module, end="")
else:
args.out.write_text(module, encoding="utf-8")
print(f"wrote {args.out}", file=sys.stderr)
return 0
if __name__ == "__main__":
raise SystemExit(main())
scripts/extract_sd_keys_and_shapes.py
-30
View File
@@ -1,30 +0,0 @@
import argparse
import json
from safetensors.torch import load_file
def extract_sd_keys_and_shapes(safetensors_file: str):
sd = load_file(safetensors_file)
keys_to_shapes = {k: v.shape for k, v in sd.items()}
out_file = "keys_and_shapes.json"
with open(out_file, "w") as f:
json.dump(keys_to_shapes, f, indent=4)
print(f"Keys and shapes written to '{out_file}'.")
def main():
parser = argparse.ArgumentParser(
description="Extracts the keys and shapes from the state dict in a safetensors file. Intended for creating "
+ "dummy state dicts for use in unit tests."
)
parser.add_argument("safetensors_file", type=str, help="Path to the safetensors file.")
args = parser.parse_args()
extract_sd_keys_and_shapes(args.safetensors_file)
if __name__ == "__main__":
main()
tests/backend/flux/redux/test_flux_redux_state_dict_utils.py
+3 -1
View File
@@ -1,6 +1,8 @@
# The state dict keys and shapes for a FLUX Redux model.
# Model source: https://huggingface.co/black-forest-labs/FLUX.1-Redux-dev/blob/1282f955f706b5240161278f2ef261d2a29ad649/flux1-redux-dev.safetensors
# The keys and shapes were extracted with extract_sd_keys_and_shapes.py.
# The keys and shapes were extracted with the predecessor of scripts/capture_state_dict_fixture.py.
# This fixture family records shapes alone; the tool writes `(shape, dtype)`, so a re-capture would
# need keys_to_mock_state_dict to read both forms.
import torch
from invokeai.backend.flux.redux.flux_redux_state_dict_utils import is_state_dict_likely_flux_redux
tests/test_capture_state_dict_fixture.py
+365
View File
@@ -0,0 +1,365 @@
"""`scripts/capture_state_dict_fixture.py` against real safetensors files.
The tool exists so a key layout can be recorded without reading a checkpoint's weights, including
one that is not on this machine. Everything it reads it reads by byte offset, so the cells here
write a real file and read it back rather than stubbing the header: an off-by-one in an offset is
the failure this tool can have, and a mocked reader would not have it.
The remote half gets the same treatment. Its cell serves the same file's bytes through an httpx
transport that honours `Range`, and asserts the two paths produce the identical module -- the point
being that the source is swappable and the parser is one.
"""
from __future__ import annotations
import importlib.util
import json
import sys
from pathlib import Path
from types import ModuleType
import httpx
import pytest
import torch
from safetensors.torch import save_file
REPO_ROOT = Path(__file__).resolve().parent.parent
MARKER = {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": 256}
HINTS = {"blocks.0.attn.qkv": {"format": "float8_e4m3fn", "full_precision_matrix_mult": True}}
def _load_module(path: Path, name: str) -> ModuleType:
"""Registered in `sys.modules` before execution: `@dataclass` resolves its own module by name."""
spec = importlib.util.spec_from_file_location(name, path)
assert spec is not None and spec.loader is not None
module = importlib.util.module_from_spec(spec)
sys.modules[name] = module
spec.loader.exec_module(module)
return module
tool = _load_module(REPO_ROOT / "scripts" / "capture_state_dict_fixture.py", "capture_state_dict_fixture")
def _blob(payload: dict) -> torch.Tensor:
return torch.frombuffer(bytearray(json.dumps(payload).encode("utf-8")), dtype=torch.uint8).clone()
@pytest.fixture
def checkpoint(tmp_path: Path) -> Path:
"""A file with two stacks of different depth and starting index, a Sequential, and a marker.
The stacks are five deep so they clear `MIN_STACK`; `refiner` starts at 2 so "keep the lowest"
is distinguishable from "keep index 0". `mlp.0`/`mlp.2` is the Sequential shape that must not be
mistaken for a stack. The marker is followed by more tensors, so a slice that reads one byte too
few or too many lands in a neighbour instead of running off the end of the file.
"""
tensors: dict[str, torch.Tensor] = {}
for index in range(5):
tensors[f"blocks.{index}.attn.qkv.weight"] = torch.zeros(6, 4, dtype=torch.int8)
tensors[f"blocks.{index}.attn.qkv.weight_scale"] = torch.zeros(6, 1, dtype=torch.float32)
tensors[f"blocks.{index}.attn.qkv.comfy_quant"] = _blob(MARKER)
tensors[f"blocks.{index}.mlp.0.weight"] = torch.zeros(4, 4, dtype=torch.bfloat16)
tensors[f"blocks.{index}.mlp.2.weight"] = torch.zeros(4, 4, dtype=torch.bfloat16)
for index in range(2, 7):
tensors[f"refiner.{index}.norm.scale"] = torch.zeros(4, dtype=torch.float32)
tensors["final.weight"] = torch.zeros(2, 4, dtype=torch.float8_e4m3fn)
tensors["final.weight_scale"] = torch.zeros((), dtype=torch.float32)
# One tensor materially larger than a marker, so "no request covers a tensor" has something to
# be wrong about: every other tensor here is small enough to slip under any sane budget.
tensors["patch_embedding.weight"] = torch.zeros(128, 128, dtype=torch.bfloat16)
path = tmp_path / "toy.safetensors"
save_file(tensors, path, metadata={"_quantization_metadata": json.dumps({"layers": HINTS})})
return path
@pytest.fixture
def range_server(checkpoint: Path, monkeypatch: pytest.MonkeyPatch):
"""Serve the real file's bytes over `Range`, recording every request the tool makes.
The recording is the point: what this tool must never do is read a tensor, and the only
evidence of that is the set of ranges it asked for.
"""
payload = checkpoint.read_bytes()
def install(status: int = 206) -> list[tuple[int, int]]:
seen: list[tuple[int, int]] = []
def serve(request: httpx.Request) -> httpx.Response:
first, last = (int(part) for part in request.headers["Range"].removeprefix("bytes=").split("-"))
seen.append((first, last))
body = payload if status == 200 else payload[first : last + 1]
return httpx.Response(status, content=body)
real_client = httpx.Client
monkeypatch.setattr(
tool.httpx, "Client", lambda **kwargs: real_client(transport=httpx.MockTransport(serve), **kwargs)
)
return seen
return install
def _undated(module: str) -> str:
"""The rendered docstring carries today's date, which differs between two captures across
midnight. Everything else about them has to be identical."""
return chr(10).join(line for line in module.splitlines() if not line.startswith("Captured on "))
def _capture(path: Path, **kwargs) -> str:
options = {"block": None, "subset": True, "want_markers": True} | kwargs
return tool.capture(str(path), **options)
def test_the_header_reports_every_key_with_its_shape_and_dtype_spelling(checkpoint: Path) -> None:
"""The dtype strings are safetensors' own, because that is what the fixtures record.
Inventing a spelling here would make every captured fixture disagree with the eleven that were
written by hand, and the disagreement would only show up in whichever test read both.
"""
source = tool.open_source(str(checkpoint))
try:
tensors, metadata = tool.read_header(source)
finally:
source.close()
assert tensors["blocks.0.attn.qkv.weight"].shape == [6, 4]
assert tensors["blocks.0.attn.qkv.weight"].dtype == "I8"
assert tensors["blocks.0.mlp.0.weight"].dtype == "BF16"
assert tensors["final.weight"].dtype == "F8_E4M3"
assert tensors["final.weight_scale"].shape == []
assert tensors["blocks.0.attn.qkv.comfy_quant"].dtype == "U8"
assert json.loads(metadata["_quantization_metadata"])["layers"] == HINTS
def test_a_marker_is_decoded_from_the_bytes_its_header_entry_points_at(checkpoint: Path) -> None:
"""Reading the blob is the whole reason the offsets are absolute rather than header-relative.
The fixtures that exist today reconstruct the marker by hand and can only assert its *length*,
which cannot tell `convrot: true` from `convrot: false`.
"""
source = tool.open_source(str(checkpoint))
try:
tensors, _metadata = tool.read_header(source)
markers = tool.read_markers(source, tensors)
finally:
source.close()
assert markers["blocks.0.attn.qkv"] == MARKER
assert set(markers) == {f"blocks.{index}.attn.qkv" for index in range(5)}
def test_the_lowest_index_of_each_stack_is_kept_and_a_sequential_is_not_one(checkpoint: Path) -> None:
"""`refiner` starts at 2, so this fails for an implementation that hard-codes index 0.
`mlp.0`/`mlp.2` are two halves of one module. Subsetting them would describe a model the file
does not contain -- and it would do so silently, since the result is still valid Python.
"""
module = _capture(checkpoint)
assert '"blocks.0.attn.qkv.weight"' in module
assert '"blocks.1.attn.qkv.weight"' not in module
assert '"refiner.2.norm.scale"' in module
assert '"refiner.3.norm.scale"' not in module
assert '"blocks.0.mlp.0.weight"' in module
assert '"blocks.0.mlp.2.weight"' in module
def test_a_requested_block_is_kept_where_the_stack_has_it(checkpoint: Path) -> None:
"""Anima's captured fixture holds block 24, not block 0, because that is where its only
`full_precision_matrix_mult` layer sits. Choosing the block is the normal case, not an escape."""
module = _capture(checkpoint, block=3)
assert '"blocks.3.attn.qkv.weight"' in module
assert '"blocks.0.attn.qkv.weight"' not in module
def test_a_stack_without_the_requested_block_keeps_its_own_lowest(checkpoint: Path, capsys) -> None:
"""One `--block` meets stacks of different depths. Refusing the capture would help nobody, so
the stack that cannot honour it says what it did instead -- on stderr and in the docstring."""
module = _capture(checkpoint, block=0)
assert '"blocks.0.attn.qkv.weight"' in module
assert '"refiner.2.norm.scale"' in module
assert "`refiner` index 2" in module
assert "has no index 0; keeping 2 instead" in capsys.readouterr().err
def test_the_rendered_module_imports_back_to_what_was_captured(checkpoint: Path, tmp_path: Path) -> None:
"""The output is a Python module the tests import, so "it is valid Python" is a property of the
product and not of the renderer -- a stray quote in a key name would only show up here."""
written = tmp_path / "captured_keys.py"
written.write_text(_capture(checkpoint), encoding="utf-8")
fixture = _load_module(written, "captured_keys")
assert fixture.state_dict_keys["blocks.0.attn.qkv.weight"] == ([6, 4], "I8")
assert fixture.state_dict_keys["final.weight_scale"] == ([], "F32")
assert fixture.markers["blocks.0.attn.qkv"] == MARKER
assert fixture.layer_hints == HINTS
def test_hints_and_markers_are_absent_when_the_file_carries_none(tmp_path: Path) -> None:
"""Krea-2's int8 build has no `__metadata__` at all. Emitting an empty `layer_hints` there would
read as "the header names no layers" rather than "there is no header block"."""
path = tmp_path / "plain.safetensors"
save_file({"final.weight": torch.zeros(2, 4, dtype=torch.bfloat16)}, path)
module = _capture(path)
assert "state_dict_keys" in module
assert "layer_hints: " not in module
assert "markers: " not in module
def test_a_hint_naming_a_dropped_block_does_not_survive_the_subset(checkpoint: Path, tmp_path: Path) -> None:
"""A header hint names a module; it is kept exactly when one of that module's tensors was.
Left in, it would name a layer the fixture no longer contains, and a test pairing hints with
weights would read that as a lost flag -- the defect this fixture family exists to catch.
"""
path = tmp_path / "later_block.safetensors"
save_file(
{key: torch.zeros(2, 2, dtype=torch.bfloat16) for key in (f"blocks.{i}.attn.qkv.weight" for i in range(5))},
path,
metadata={"_quantization_metadata": json.dumps({"layers": {"blocks.4.attn.qkv": {"format": "x"}}})},
)
assert "layer_hints: " not in _capture(path)
assert "layer_hints: " in _capture(path, block=4)
def test_the_remote_source_captures_the_same_module_as_the_local_one(checkpoint: Path, range_server) -> None:
"""The remote half is the one that can reach a layout nobody here owns, so it is the half that
must not drift. Serving the same bytes through `Range` has to produce the identical module."""
seen = range_server()
remote = tool.capture("https://example.invalid/toy.safetensors", block=None, subset=True, want_markers=True)
assert seen[0] == (0, 7)
assert _undated(remote.replace("https://example.invalid/toy.safetensors", checkpoint.as_posix())) == _undated(
_capture(checkpoint)
)
def test_a_source_name_cannot_close_the_docstring_it_lands_in(range_server) -> None:
"""The source name is written into the generated module's docstring, so it is code, not text.
A local path is rendered with posix separators and so cannot carry an escape, but a URL is
whatever the operator passed. Driven through the remote source because that is the one whose
name the tool does not normalise.
"""
range_server()
module = tool.capture(r'https://example.invalid/a"""b\c.safetensors', block=None, subset=True, want_markers=False)
compile(module, "<captured>", "exec")
def test_a_bare_repo_path_resolves_against_hugging_face(tmp_path: Path) -> None:
"""`<org>/<repo>/<path>` is how the checkpoints are named everywhere else in this work. An
existing local path wins, so a relative filename is never mistaken for a repository."""
argument = "Comfy-Org/Wan_2.2_Repackaged/split_files/diffusion_models/wan.safetensors"
source = tool.open_source(argument)
try:
assert (
source._url
== f"https://huggingface.co/Comfy-Org/Wan_2.2_Repackaged/resolve/main/{argument.split('/', 2)[2]}"
)
# The docstring cites the repository, which is what a reader needs to find the file again.
assert str(source) == argument
finally:
source.close()
local = tmp_path / "local.safetensors"
save_file({"w": torch.zeros(2)}, local)
source = tool.open_source(str(local))
try:
assert isinstance(source, tool.LocalFile)
finally:
source.close()
def test_only_the_markers_that_survived_the_subset_are_read(checkpoint: Path, tmp_path: Path, range_server) -> None:
"""Reading markers for dropped blocks costs a request each and emits layers the fixture does
not contain.
The toy file has one marker per block and five blocks; a capture that kept block 0 must ask for
exactly one. On a 40-block checkpoint the same slip is 40 round trips and a `markers` mapping
whose keys no `state_dict_keys` entry backs -- which is what the consumer cells assert against,
so it would be found only after the bad fixture was already checked in.
"""
seen = range_server()
written = tmp_path / "subset_markers.py"
written.write_text(
tool.capture("https://example.invalid/toy.safetensors", block=None, subset=True, want_markers=True),
encoding="utf-8",
)
assert set(_load_module(written, "subset_markers").markers) == {"blocks.0.attn.qkv"}
assert len(seen) == 3 # the eight length bytes, the header, one marker
def test_no_request_ever_covers_a_tensor(checkpoint: Path, range_server) -> None:
"""The whole promise of the tool: a capture of a multi-gigabyte file costs kilobytes.
Asserted as a property of the requests rather than of the runtime, so it holds for a file this
suite is too small to contain. The header read is exempt and named; everything else is a marker.
"""
seen = range_server()
tool.capture("https://example.invalid/toy.safetensors", block=None, subset=True, want_markers=True)
lengths = [last - first + 1 for first, last in seen[2:]]
assert lengths
assert max(lengths) < 1024, seen
def test_a_server_that_ignores_the_range_is_refused(range_server) -> None:
"""A 200 means the body is the whole file, already buffered by the time we could look at it.
That is the one cost this tool exists to avoid, and it is invisible otherwise: the capture
succeeds, it just downloaded 14 GB to read four kilobytes of it.
"""
range_server(status=200)
with pytest.raises(SystemExit, match="does not honour"):
tool.capture("https://example.invalid/toy.safetensors", block=None, subset=True, want_markers=True)
def test_something_that_is_not_a_safetensors_is_refused_before_the_second_read(tmp_path: Path, range_server) -> None:
"""The header length comes out of the file, so an HTML error page or an LFS pointer yields a
nonsense one -- and the next read asks for exactly that many bytes."""
not_a_checkpoint = tmp_path / "index.html"
not_a_checkpoint.write_bytes(b"<!DOCTYPE html><title>404</title>" + b"x" * 64)
with pytest.raises(SystemExit, match="not a safetensors file"):
tool.capture(str(not_a_checkpoint), block=None, subset=True, want_markers=False)
def test_a_mistyped_local_path_is_not_fetched_from_hugging_face() -> None:
"""`S:/ModelStuff3/typo.safetensors` has three segments and does not exist, so without a guard
it becomes `huggingface.co/S:/ModelStuff3/resolve/main/typo.safetensors` and the operator is
told about an HTTP status instead of a missing file."""
for mistyped in (r"S:/ModelStuff3/typo.safetensors", "/home/me/models/typo.safetensors"):
with pytest.raises(SystemExit, match="no such file"):
tool.open_source(mistyped)
def test_a_non_finite_header_value_is_refused_rather_than_rendered(tmp_path: Path) -> None:
"""Python's JSON accepts `NaN` and `Infinity`; Python's `repr` of them is a bare name.
Rendered into a fixture that is a `NameError` at import, i.e. a module that looks captured and
cannot be read. No producer writes them, so this is about what happens when one does.
"""
path = tmp_path / "nan.safetensors"
save_file(
{"final.weight": torch.zeros(2, 2, dtype=torch.bfloat16)},
path,
metadata={"_quantization_metadata": '{"layers": {"final": {"scale": NaN}}}'},
)
with pytest.raises(SystemExit, match="non-finite"):
_capture(path)