SD15 Safetensor VAE Fix

# sd15 repairs checkpoints with vae errors
# based on https://github.com/huggingface/diffusers/issues/4790
# tested on https://civitai.com/models/146136
# use with Easy Diffusion:
# 1. Open the Developer Console
# 2. python /path/to/script.py -i /path/to/checkpoint.safetensors -o /path/to/fixed_checkpoint.safetensors

import argparse
from safetensors import safe_open
from safetensors.torch import save_file
from pathlib import Path

def fix_diffusers_model_conversion(load_path, save_path):
    # load original
    tensors = {}
    with safe_open(load_path, framework="pt") as f:
        for key in f.keys():
            tensors[key] = f.get_tensor(key)
    # migrate
    new_tensors = {}
    for k, v in tensors.items():
        new_key = k
        # only fix the vae
        if 'first_stage_model.' in k:
            # migrate q, k, v keys
            new_key = new_key.replace('.to_q.weight', '.q.weight')
            new_key = new_key.replace('.to_q.bias', '.q.bias')
            new_key = new_key.replace('.to_k.weight', '.k.weight')
            new_key = new_key.replace('.to_k.bias', '.k.bias')
            new_key = new_key.replace('.to_v.weight', '.v.weight')
            new_key = new_key.replace('.to_v.bias', '.v.bias')
        new_tensors[new_key] = v
    # save
    save_file(new_tensors, save_path)

# get the cli arguments
parser = argparse.ArgumentParser()
parser.add_argument('-i', '--input', required=True, help='Broken checkpoint path')
parser.add_argument('-o', '--output', required=True, help='Path for newly fixed checkpoint')
args = parser.parse_args()

# check the input file exists and the output file does not
input_file = Path(args.input)
output_file = Path(args.output)

if input_file.is_file() and not output_file.is_file():
    print("Attempting to fix checkpoint...")
    fix_diffusers_model_conversion(args.input, args.output)
else:
    print("Invalid command line. Use -h for help.")