File size: 3,286 Bytes
eae424a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | #!/usr/bin/env python3
"""Convert DinoVision's runtime decoder.bin into named SafeTensors."""
from __future__ import annotations
import argparse
import hashlib
import json
from pathlib import Path
import numpy as np
from safetensors.numpy import load_file, save_file
STAGES = (256, 128, 64, 32)
BLEND_STAGES = 2
def sha256(path: Path) -> str:
hasher = hashlib.sha256()
with path.open("rb") as source:
for chunk in iter(lambda: source.read(1024 * 1024), b""):
hasher.update(chunk)
return hasher.hexdigest()
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--training", required=True, type=Path)
parser.add_argument("--input", type=Path)
parser.add_argument("--output", required=True, type=Path)
args = parser.parse_args()
training = json.loads(args.training.read_text(encoding="utf-8"))
source = args.input or args.training.parent / "decoder.bin"
source_hash = sha256(source)
if source_hash != training["decoder_sha256"]:
raise SystemExit(
f"{source} has SHA-256 {source_hash}; training record expects "
f"{training['decoder_sha256']}"
)
flat = np.fromfile(source, dtype="<f4")
offset = 0
tensors: dict[str, np.ndarray] = {}
def take(name: str, shape: tuple[int, ...]) -> None:
nonlocal offset
count = int(np.prod(shape))
end = offset + count
if end > len(flat):
raise SystemExit(f"{source} ends inside tensor {name}")
tensors[name] = flat[offset:end].reshape(shape).copy()
offset = end
def block(name: str, in_channels: int, out_channels: int) -> None:
take(f"{name}.weight", (out_channels, in_channels, 3, 3))
take(f"{name}.bias", (out_channels,))
take(f"{name}.norm.weight", (out_channels,))
take(f"{name}.norm.bias", (out_channels,))
in_channels = 384
for index, out_channels in enumerate(STAGES):
block(f"dec.{index}", in_channels, out_channels)
if index < BLEND_STAGES:
block(f"dec.{index}b", out_channels, out_channels)
in_channels = out_channels
take("dec.out.weight", (3, in_channels, 3, 3))
take("dec.out.bias", (3,))
if offset != len(flat):
raise SystemExit(f"{source} has {len(flat) - offset} trailing f32 values")
metadata = {
"format": "dinovision-decoder-v1",
"source_decoder_sha256": source_hash,
"model_sha256": training["model_sha256"],
"dataset_manifest_sha256": training["dataset_manifest_sha256"],
"encoder_layers": str(training["encoder_layers"]),
"image_size": str(training["image_size"]),
"seed": str(training["seed"]),
"objective": training["objective"],
}
args.output.parent.mkdir(parents=True, exist_ok=True)
save_file(tensors, args.output, metadata=metadata)
loaded = load_file(args.output)
for name, expected in tensors.items():
if name not in loaded or not np.array_equal(expected, loaded[name]):
raise SystemExit(f"SafeTensors roundtrip changed tensor {name}")
print(
f"wrote {args.output}: {len(tensors)} named tensors, "
f"SHA-256 {sha256(args.output)}"
)
if __name__ == "__main__":
main()
|