| |
| """Small helpers for prepare_model_dir.sh (standard library only). |
| |
| Run as: python3 -I prepare_helpers.py <command> ... |
| Commands: vae-config, rewrite-source, st-check, shards-check, verify |
| """ |
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| import os |
| import re |
| import struct |
| import sys |
| import tempfile |
| from pathlib import Path |
|
|
| ENCODER_SOURCE = "hf://FastVideo/FastH3-Omni-NVFP4-QAD/text_encoder" |
| MANIFEST = "mlx_h3_nvfp4_encoder.json" |
|
|
|
|
| def _atomic_write(path: Path, text: str) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| fd, tmp = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent) |
| with os.fdopen(fd, "w") as f: |
| f.write(text) |
| os.replace(tmp, path) |
|
|
|
|
| def safetensors_header(path: Path) -> dict: |
| """Parse the header and check the file size matches the declared tensor data.""" |
| size = path.stat().st_size |
| with path.open("rb") as f: |
| raw = f.read(8) |
| if len(raw) != 8: |
| raise ValueError("file shorter than 8 bytes") |
| (n,) = struct.unpack("<Q", raw) |
| if n <= 0 or 8 + n > size: |
| raise ValueError(f"bad header length {n} for file size {size}") |
| header = json.loads(f.read(n)) |
| end = max((v["data_offsets"][1] for k, v in header.items() if k != "__metadata__"), default=0) |
| if 8 + n + end != size: |
| raise ValueError(f"truncated or padded: expected {8 + n + end} bytes, found {size}") |
| return header |
|
|
|
|
| def cmd_vae_config(src: str, dst: str) -> None: |
| meta = safetensors_header(Path(src)).get("__metadata__") or {} |
| config = json.loads(meta["architecture"])["teacher_config"] |
| _atomic_write(Path(dst), json.dumps(config, indent=2)) |
|
|
|
|
| def cmd_rewrite_source(manifest: str, source: str) -> None: |
| path = Path(manifest) |
| data = json.loads(path.read_text()) |
| if data.get("source_dir") != source: |
| data["source_dir"] = source |
| _atomic_write(path, json.dumps(data, indent=2) + "\n") |
|
|
|
|
| def cmd_st_check(files: list[str]) -> int: |
| bad = 0 |
| for name in files: |
| try: |
| safetensors_header(Path(name)) |
| except Exception as exc: |
| print(f"BAD safetensors {name}: {exc}", file=sys.stderr) |
| bad += 1 |
| return 1 if bad else 0 |
|
|
|
|
| def cmd_shards_check(directory: str, prefix: str) -> int: |
| """All `<prefix>-0000k-of-0000N.safetensors` shards present and well-formed.""" |
| d = Path(directory) |
| pat = re.compile(rf"^{re.escape(prefix)}-(\d+)-of-(\d+)\.safetensors$") |
| found = {} |
| total = 0 |
| for p in d.iterdir() if d.is_dir() else []: |
| m = pat.match(p.name) |
| if m: |
| found[int(m.group(1))] = p |
| total = int(m.group(2)) |
| if not found or sorted(found) != list(range(1, total + 1)): |
| print(f"shards incomplete in {d}: have {sorted(found)} of {total}", file=sys.stderr) |
| return 1 |
| return cmd_st_check([str(p) for p in found.values()]) |
|
|
|
|
| def sha256_file(path: Path) -> str: |
| h = hashlib.sha256() |
| with path.open("rb") as f: |
| while chunk := f.read(16 << 20): |
| h.update(chunk) |
| return h.hexdigest() |
|
|
|
|
| def _is_converter_output(rel: str) -> bool: |
| return rel.startswith("dit/") or rel in (f"root/text_encoder/{MANIFEST}", "root/text_encoder/model.safetensors") |
|
|
|
|
| def cmd_verify(model_dir: str, expected_file: str, formats: list[str]) -> int: |
| root = Path(model_dir) |
| fails, warns, ok = [], [], 0 |
| for line in Path(expected_file).read_text().splitlines(): |
| if not line.strip(): |
| continue |
| digest, rel = line.split(None, 1) |
| rel = rel.strip() |
| if rel.startswith("dit/") and rel.split("/")[1] not in formats: |
| continue |
| path = root / rel |
| soft = _is_converter_output(rel) |
| if not path.is_file(): |
| fails.append(f"missing: {rel}") |
| print(f" MISSING {rel}", flush=True) |
| continue |
| got = sha256_file(path) |
| if got == digest: |
| ok += 1 |
| print(f" ok {rel}", flush=True) |
| elif soft: |
| warns.append(f"converter output differs (not bit-identical): {rel}") |
| print(f" WARN {rel} (converter output differs from reference)", flush=True) |
| else: |
| fails.append(f"hash mismatch: {rel}") |
| print(f" FAIL {rel}\n expected {digest}\n got {got}", flush=True) |
| |
| for fmt in formats: |
| d = root / "dit" / fmt |
| for name in ("mlx_h3_dit.json", "mlx_h3_dit.safetensors"): |
| if not (d / name).is_file(): |
| fails.append(f"missing: dit/{fmt}/{name}") |
| if (d / "mlx_h3_dit.safetensors").is_file() and cmd_st_check([str(d / "mlx_h3_dit.safetensors")]): |
| fails.append(f"dit/{fmt}/mlx_h3_dit.safetensors is not a complete safetensors file") |
| enc = root / "root/text_encoder/model.safetensors" |
| if enc.is_file() and cmd_st_check([str(enc)]): |
| fails.append("root/text_encoder/model.safetensors is not a complete safetensors file") |
| for rel in ("root/transformer_ref/config.json", "root/vae/config.json", "root/fastvideo_inference.json"): |
| try: |
| json.loads((root / rel).read_text()) |
| except Exception as exc: |
| fails.append(f"unreadable JSON {rel}: {exc}") |
| print(f"\nverify: {ok} exact matches, {len(warns)} warnings, {len(fails)} failures") |
| for w in warns: |
| print(f" warning: {w}") |
| for f in fails: |
| print(f" FAILURE: {f}") |
| return 1 if fails else 0 |
|
|
|
|
| def main(argv: list[str]) -> int: |
| if not argv: |
| print(__doc__) |
| return 2 |
| cmd, args = argv[0], argv[1:] |
| if cmd == "vae-config": |
| cmd_vae_config(*args) |
| elif cmd == "rewrite-source": |
| cmd_rewrite_source(*args) |
| elif cmd == "st-check": |
| return cmd_st_check(args) |
| elif cmd == "shards-check": |
| return cmd_shards_check(*args) |
| elif cmd == "verify": |
| return cmd_verify(args[0], args[1], args[2:]) |
| else: |
| print(f"unknown command {cmd}", file=sys.stderr) |
| return 2 |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main(sys.argv[1:])) |
|
|
