aichan_blend / safetensors_converter.py
Kaeya's picture
Update safetensors_converter.py
cd5a91a
raw
history blame contribute delete
No virus
1.72 kB
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)