| |
| """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() |
|
|
