File size: 1,718 Bytes
9d71531 9153408 9d71531 cd5a91a 9153408 9d71531 9153408 9d71531 9153408 9d71531 9153408 9d71531 9153408 9d71531 9153408 |
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 |
import argparse
from pathlib import Path
import torch
from safetensors.torch import save_file
def convert(path: Path, half: bool = False, no_ema: bool = False):
state_dict = torch.load(path, map_location="cpu")
if "state_dict" in state_dict:
state_dict = state_dict["state_dict"]
to_remove = []
for k, v in state_dict.items():
if not isinstance(v, torch.Tensor):
to_remove.append(k)
elif no_ema and "ema" in k:
to_remove.append(k)
for k in to_remove:
del state_dict[k]
if half:
state_dict = {k: v.half() for k, v in state_dict.items()}
output_name = path.stem
if no_ema:
output_name += "-pruned"
if half:
output_name += "-fp16"
output_path = path.parent / f"{output_name}.safetensors"
save_file(state_dict, output_path.as_posix())
def main(path: str, half: bool = False, no_ema: bool = False):
path_ = Path(path).resolve()
if not path_.exists():
raise ValueError(f"Invalid path: {path}")
if path_.is_file():
to_convert = [path_]
else:
to_convert = list(path_.glob("*.ckpt"))
for file in to_convert:
print(f"Converting... {file}")
convert(file, half, no_ema)
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("path", type=str, help="Path to checkpoint file or directory.")
parser.add_argument(
"--half", action="store_true", help="Convert to half precision."
)
parser.add_argument("--no-ema", action="store_true", help="Ignore EMA weights.")
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
main(args.path, args.half, args.no_ema)
|