Spaces:
Sleeping
Sleeping
| """Convert the owner's checkpoint without enabling unrestricted pickle loading.""" | |
| import argparse | |
| from pathlib import Path | |
| def convert(checkpoint, output): | |
| import torch | |
| from safetensors.torch import save_file | |
| import os | |
| import tempfile | |
| checkpoint, output = Path(checkpoint), Path(output) | |
| if output.exists(): | |
| raise ValueError("Output already exists; choose another path.") | |
| obj = torch.load(checkpoint, map_location="cpu", weights_only=True) | |
| state = obj.get("state_dict", obj) if isinstance(obj, dict) else None | |
| if not isinstance(state, dict) or not state: | |
| raise ValueError("No model tensors were found.") | |
| if any(not isinstance(k, str) or not isinstance(v, torch.Tensor) for k, v in state.items()): | |
| raise ValueError("Unexpected checkpoint content. Inspect it before conversion.") | |
| fd, temp_name = tempfile.mkstemp(dir=output.parent, suffix=".safetensors") | |
| os.close(fd) | |
| try: | |
| save_file({k: v.detach().contiguous().clone() for k, v in state.items()}, temp_name) | |
| os.replace(temp_name, output) | |
| finally: | |
| Path(temp_name).unlink(missing_ok=True) | |
| return output | |
| def main(): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("checkpoint", type=Path) | |
| parser.add_argument("output", type=Path) | |
| args = parser.parse_args() | |
| convert(args.checkpoint, args.output) | |
| print("Saved", args.output) | |
| if __name__ == "__main__": | |
| main() | |