Initial release
This commit is contained in:
civ
@@ -0,0 +1,335 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Merge selected tensors from multiple safetensors files without loading weights.
|
||||
|
||||
Edit ``OUTPUT_PATH`` and ``SOURCE_RULES`` below, then run:
|
||||
|
||||
python scripts/merge_safetensors.py
|
||||
|
||||
Each source rule uses regular expressions against complete tensor names.
|
||||
``include`` is required and matches when any expression succeeds. ``exclude``
|
||||
wins over ``include``. Expressions are evaluated with ``re.search``.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import struct
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import BinaryIO
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Configuration
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
OUTPUT_PATH = Path(r"..\models\diffusion_models\minimax_h3_ref2va_pruned_bf16.safetensors")
|
||||
|
||||
SOURCE_RULES = [
|
||||
{
|
||||
"path": Path(r"..\models\diffusion_models\minimax_h3_ref2va_bf16.safetensors"),
|
||||
"include": [r".*"],
|
||||
"exclude": [r".*adaln_proj\.linear.*", r"time_embedder.*"],
|
||||
},
|
||||
{
|
||||
"path": Path(r"..\models\diffusion_models\minimax_h3_ref2va_pruned_int8_convrot.safetensors"),
|
||||
"include": [r"^.*adaln_proj\.linear.*", "adaln_t_table"],
|
||||
"exclude": [],
|
||||
},
|
||||
]
|
||||
|
||||
# Safetensors metadata is optional. Set this to a dict[str, str] if needed.
|
||||
OUTPUT_METADATA = None
|
||||
|
||||
# Refuse to replace an existing output unless explicitly enabled.
|
||||
OVERWRITE_OUTPUT = False
|
||||
|
||||
# Only tensor headers and this fixed-size buffer are held in memory.
|
||||
COPY_BUFFER_SIZE = 8 * 1024 * 1024
|
||||
PROGRESS_INTERVAL = 1024 * 1024 * 1024
|
||||
MAX_HEADER_SIZE = 256 * 1024 * 1024
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TensorEntry:
|
||||
name: str
|
||||
source_path: Path
|
||||
source_data_offset: int
|
||||
source_start: int
|
||||
source_end: int
|
||||
dtype: str
|
||||
shape: list[int]
|
||||
|
||||
@property
|
||||
def size(self) -> int:
|
||||
return self.source_end - self.source_start
|
||||
|
||||
|
||||
def format_bytes(size: int) -> str:
|
||||
value = float(size)
|
||||
for unit in ("B", "KiB", "MiB", "GiB", "TiB"):
|
||||
if value < 1024.0 or unit == "TiB":
|
||||
return f"{value:.2f} {unit}"
|
||||
value /= 1024.0
|
||||
raise AssertionError("unreachable")
|
||||
|
||||
|
||||
def read_exact(file: BinaryIO, size: int, description: str) -> bytes:
|
||||
data = file.read(size)
|
||||
if len(data) != size:
|
||||
raise ValueError(f"truncated {description}: expected {size} bytes, got {len(data)}")
|
||||
return data
|
||||
|
||||
|
||||
def read_safetensors_header(path: Path) -> tuple[dict, int, int]:
|
||||
file_size = path.stat().st_size
|
||||
with path.open("rb") as file:
|
||||
header_size = struct.unpack("<Q", read_exact(file, 8, f"header size in {path}"))[0]
|
||||
if header_size == 0 or header_size > MAX_HEADER_SIZE:
|
||||
raise ValueError(
|
||||
f"invalid header size in {path}: {header_size} "
|
||||
f"(limit: {MAX_HEADER_SIZE})"
|
||||
)
|
||||
header_bytes = read_exact(file, header_size, f"header in {path}")
|
||||
|
||||
try:
|
||||
header = json.loads(header_bytes)
|
||||
except (UnicodeDecodeError, json.JSONDecodeError) as error:
|
||||
raise ValueError(f"invalid safetensors JSON header in {path}: {error}") from error
|
||||
if not isinstance(header, dict):
|
||||
raise ValueError(f"safetensors header in {path} is not an object")
|
||||
|
||||
data_offset = 8 + header_size
|
||||
if data_offset > file_size:
|
||||
raise ValueError(f"safetensors data offset is past end of file: {path}")
|
||||
return header, data_offset, file_size
|
||||
|
||||
|
||||
def parse_tensor_entry(
|
||||
name: str,
|
||||
info: object,
|
||||
source_path: Path,
|
||||
source_data_offset: int,
|
||||
source_file_size: int,
|
||||
) -> TensorEntry:
|
||||
if not isinstance(info, dict):
|
||||
raise ValueError(f"{source_path}: tensor {name!r} has an invalid header entry")
|
||||
|
||||
dtype = info.get("dtype")
|
||||
shape = info.get("shape")
|
||||
offsets = info.get("data_offsets")
|
||||
if not isinstance(dtype, str):
|
||||
raise ValueError(f"{source_path}: tensor {name!r} has an invalid dtype")
|
||||
if not isinstance(shape, list) or not all(
|
||||
isinstance(dimension, int) and dimension >= 0 for dimension in shape
|
||||
):
|
||||
raise ValueError(f"{source_path}: tensor {name!r} has an invalid shape")
|
||||
if (
|
||||
not isinstance(offsets, list)
|
||||
or len(offsets) != 2
|
||||
or not all(isinstance(offset, int) for offset in offsets)
|
||||
):
|
||||
raise ValueError(f"{source_path}: tensor {name!r} has invalid data offsets")
|
||||
|
||||
start, end = offsets
|
||||
if start < 0 or end < start or source_data_offset + end > source_file_size:
|
||||
raise ValueError(
|
||||
f"{source_path}: tensor {name!r} byte range [{start}, {end}) "
|
||||
"is outside the file"
|
||||
)
|
||||
|
||||
return TensorEntry(
|
||||
name=name,
|
||||
source_path=source_path,
|
||||
source_data_offset=source_data_offset,
|
||||
source_start=start,
|
||||
source_end=end,
|
||||
dtype=dtype,
|
||||
shape=list(shape),
|
||||
)
|
||||
|
||||
|
||||
def compile_patterns(rule_index: int, field: str, values: object) -> list[re.Pattern[str]]:
|
||||
if not isinstance(values, list) or not all(isinstance(value, str) for value in values):
|
||||
raise TypeError(f"SOURCE_RULES[{rule_index}][{field!r}] must be a list of strings")
|
||||
try:
|
||||
return [re.compile(value) for value in values]
|
||||
except re.error as error:
|
||||
raise ValueError(
|
||||
f"invalid regex in SOURCE_RULES[{rule_index}][{field!r}]: {error}"
|
||||
) from error
|
||||
|
||||
|
||||
def collect_entries() -> list[TensorEntry]:
|
||||
if not SOURCE_RULES:
|
||||
raise ValueError("SOURCE_RULES must contain at least one source")
|
||||
|
||||
entries: list[TensorEntry] = []
|
||||
selected_by_name: dict[str, TensorEntry] = {}
|
||||
header_cache: dict[Path, tuple[dict, int, int]] = {}
|
||||
|
||||
for rule_index, rule in enumerate(SOURCE_RULES):
|
||||
if not isinstance(rule, dict) or "path" not in rule or "include" not in rule:
|
||||
raise TypeError(
|
||||
f"SOURCE_RULES[{rule_index}] must contain 'path' and 'include'"
|
||||
)
|
||||
|
||||
source_path = Path(rule["path"])
|
||||
if not source_path.is_file():
|
||||
raise FileNotFoundError(f"source file does not exist: {source_path}")
|
||||
source_path = source_path.resolve()
|
||||
|
||||
include = compile_patterns(rule_index, "include", rule["include"])
|
||||
exclude = compile_patterns(rule_index, "exclude", rule.get("exclude", []))
|
||||
if not include:
|
||||
raise ValueError(f"SOURCE_RULES[{rule_index}]['include'] must not be empty")
|
||||
|
||||
if source_path not in header_cache:
|
||||
header_cache[source_path] = read_safetensors_header(source_path)
|
||||
header, data_offset, file_size = header_cache[source_path]
|
||||
|
||||
matched = 0
|
||||
for name, info in header.items():
|
||||
if name == "__metadata__":
|
||||
continue
|
||||
if not any(pattern.search(name) for pattern in include):
|
||||
continue
|
||||
if any(pattern.search(name) for pattern in exclude):
|
||||
continue
|
||||
|
||||
entry = parse_tensor_entry(name, info, source_path, data_offset, file_size)
|
||||
previous = selected_by_name.get(name)
|
||||
if previous is not None:
|
||||
raise ValueError(
|
||||
f"tensor {name!r} was selected more than once:\n"
|
||||
f" first: {previous.source_path}\n"
|
||||
f" second: {source_path}"
|
||||
)
|
||||
selected_by_name[name] = entry
|
||||
print(f"entry {entry}")
|
||||
entries.append(entry)
|
||||
matched += 1
|
||||
|
||||
print(f"Rule {rule_index}: selected {matched} tensors from {source_path}")
|
||||
if matched == 0:
|
||||
raise ValueError(
|
||||
f"SOURCE_RULES[{rule_index}] did not select any tensors; check its regexes"
|
||||
)
|
||||
|
||||
if not entries:
|
||||
raise ValueError("no tensors were selected")
|
||||
return entries
|
||||
|
||||
|
||||
def build_output_header(entries: list[TensorEntry]) -> tuple[bytes, int]:
|
||||
header: dict[str, object] = {}
|
||||
if OUTPUT_METADATA is not None:
|
||||
if not isinstance(OUTPUT_METADATA, dict) or not all(
|
||||
isinstance(key, str) and isinstance(value, str)
|
||||
for key, value in OUTPUT_METADATA.items()
|
||||
):
|
||||
raise TypeError("OUTPUT_METADATA must be None or a dict[str, str]")
|
||||
header["__metadata__"] = OUTPUT_METADATA
|
||||
|
||||
output_offset = 0
|
||||
for entry in entries:
|
||||
header[entry.name] = {
|
||||
"dtype": entry.dtype,
|
||||
"shape": entry.shape,
|
||||
"data_offsets": [output_offset, output_offset + entry.size],
|
||||
}
|
||||
output_offset += entry.size
|
||||
|
||||
header_bytes = json.dumps(header, separators=(",", ":"), ensure_ascii=False).encode(
|
||||
"utf-8"
|
||||
)
|
||||
header_bytes += b" " * (-len(header_bytes) % 8)
|
||||
return header_bytes, output_offset
|
||||
|
||||
|
||||
def copy_tensor(source: BinaryIO, output: BinaryIO, entry: TensorEntry) -> None:
|
||||
source.seek(entry.source_data_offset + entry.source_start)
|
||||
remaining = entry.size
|
||||
while remaining:
|
||||
chunk = source.read(min(COPY_BUFFER_SIZE, remaining))
|
||||
if not chunk:
|
||||
raise OSError(
|
||||
f"unexpected end of file while copying {entry.name!r} "
|
||||
f"from {entry.source_path}"
|
||||
)
|
||||
output.write(chunk)
|
||||
remaining -= len(chunk)
|
||||
|
||||
|
||||
def write_output(entries: list[TensorEntry]) -> None:
|
||||
if COPY_BUFFER_SIZE <= 0:
|
||||
raise ValueError("COPY_BUFFER_SIZE must be positive")
|
||||
|
||||
output_path = OUTPUT_PATH.resolve()
|
||||
source_paths = {entry.source_path.resolve() for entry in entries}
|
||||
if output_path in source_paths:
|
||||
raise ValueError("OUTPUT_PATH must not be one of the source files")
|
||||
if output_path.exists() and not OVERWRITE_OUTPUT:
|
||||
raise FileExistsError(
|
||||
f"output already exists: {output_path}; set OVERWRITE_OUTPUT = True to replace it"
|
||||
)
|
||||
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
partial_path = output_path.with_name(output_path.name + ".partial")
|
||||
if partial_path.exists():
|
||||
raise FileExistsError(
|
||||
f"partial output already exists: {partial_path}; remove it before retrying"
|
||||
)
|
||||
|
||||
header_bytes, tensor_bytes = build_output_header(entries)
|
||||
print(
|
||||
f"Writing {len(entries)} tensors ({format_bytes(tensor_bytes)}) to {output_path}"
|
||||
)
|
||||
|
||||
current_source_path: Path | None = None
|
||||
current_source: BinaryIO | None = None
|
||||
copied = 0
|
||||
next_progress = PROGRESS_INTERVAL
|
||||
try:
|
||||
with partial_path.open("xb") as output:
|
||||
output.write(struct.pack("<Q", len(header_bytes)))
|
||||
output.write(header_bytes)
|
||||
|
||||
try:
|
||||
for entry in entries:
|
||||
if entry.source_path != current_source_path:
|
||||
if current_source is not None:
|
||||
current_source.close()
|
||||
current_source = entry.source_path.open("rb")
|
||||
current_source_path = entry.source_path
|
||||
|
||||
copy_tensor(current_source, output, entry)
|
||||
copied += entry.size
|
||||
if PROGRESS_INTERVAL > 0 and copied >= next_progress:
|
||||
print(
|
||||
f" copied {format_bytes(copied)} / "
|
||||
f"{format_bytes(tensor_bytes)}"
|
||||
)
|
||||
while next_progress <= copied:
|
||||
next_progress += PROGRESS_INTERVAL
|
||||
finally:
|
||||
if current_source is not None:
|
||||
current_source.close()
|
||||
|
||||
if copied != tensor_bytes:
|
||||
raise OSError(f"copied {copied} tensor bytes, expected {tensor_bytes}")
|
||||
os.replace(partial_path, output_path)
|
||||
except BaseException:
|
||||
partial_path.unlink(missing_ok=True)
|
||||
raise
|
||||
|
||||
print(f"Done: {output_path} ({format_bytes(output_path.stat().st_size)})")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
entries = collect_entries()
|
||||
write_output(entries)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user