Spaces:
Runtime error
Runtime error
File size: 1,116 Bytes
c19ca42 |
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 |
#!/usr/bin/env python3
"""Extracts the configuration file from a slim inference checkpoint."""
import argparse
import json
from pathlib import Path
import sys
import k_diffusion as K
import safetensors.torch as safetorch
def main():
p = argparse.ArgumentParser(description=__doc__,
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
p.add_argument("checkpoint", type=Path,
help="the inference checkpoint to extract the configuration from")
p.add_argument("--output", "-o", type=Path,
help="the output configuration file")
args = p.parse_args()
print(f"Loading inference checkpoint {args.checkpoint}...", file=sys.stderr)
metadata = K.utils.get_safetensors_metadata(args.checkpoint)
if "config" not in metadata:
raise ValueError("No configuration found in checkpoint")
output_path = args.output or args.checkpoint.with_suffix(".json")
print(f"Saving configuration to {output_path}...", file=sys.stderr)
output_path.write_text(metadata["config"])
if __name__ == "__main__":
main()
|