From 2edfb3e2dfacfaf94d0571162b18f132e59c2231 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sat, 26 Sep 2026 04:40:23 +0200 Subject: [PATCH] test(fixtures): capture a key layout from the safetensors header alone 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. --- scripts/capture_state_dict_fixture.py | 413 ++++++++++++++++++ scripts/extract_sd_keys_and_shapes.py | 30 -- .../redux/test_flux_redux_state_dict_utils.py | 4 +- tests/test_capture_state_dict_fixture.py | 365 ++++++++++++++++ 4 files changed, 781 insertions(+), 31 deletions(-) create mode 100644 scripts/capture_state_dict_fixture.py delete mode 100644 scripts/extract_sd_keys_and_shapes.py create mode 100644 tests/test_capture_state_dict_fixture.py diff --git a/scripts/capture_state_dict_fixture.py b/scripts/capture_state_dict_fixture.py new file mode 100644 index 0000000000..8184f93521 --- /dev/null +++ b/scripts/capture_state_dict_fixture.py @@ -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 --out + +`` is a local path, a full `https://...` URL, or `//`, 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 `//` 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 //") + 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(" dict[str, dict[str, Any]]: + """Decode every `.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 ''} 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 ''}` 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 //") + 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()) diff --git a/scripts/extract_sd_keys_and_shapes.py b/scripts/extract_sd_keys_and_shapes.py deleted file mode 100644 index 33c7e83627..0000000000 --- a/scripts/extract_sd_keys_and_shapes.py +++ /dev/null @@ -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() diff --git a/tests/backend/flux/redux/test_flux_redux_state_dict_utils.py b/tests/backend/flux/redux/test_flux_redux_state_dict_utils.py index 733ac9f4b7..e3e056e15b 100644 --- a/tests/backend/flux/redux/test_flux_redux_state_dict_utils.py +++ b/tests/backend/flux/redux/test_flux_redux_state_dict_utils.py @@ -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 diff --git a/tests/test_capture_state_dict_fixture.py b/tests/test_capture_state_dict_fixture.py new file mode 100644 index 0000000000..9f81d2f012 --- /dev/null +++ b/tests/test_capture_state_dict_fixture.py @@ -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, "", "exec") + + +def test_a_bare_repo_path_resolves_against_hugging_face(tmp_path: Path) -> None: + """`//` 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"404" + 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)