# Got a bunch of .ckpt files to convert?
# Here's a handy script to take care of all that for you!
# Original .ckpt files are not touched!
# Make sure you have enough disk space! You are going to DOUBLE the size of your models folder!
#
# First, run:
# pip install torch torchsde==0.2.5 safetensors==0.2.5
#
# Place this file in the **SAME DIRECTORY** as all of your .ckpt files, open a command prompt for that folder, and run:
# python convert_to_safe.py

import os
import torch
from safetensors.torch import save_file

files = os.listdir()
for f in files:
    if f.lower().endswith('.bin'):
        print(f'Loading {f}...')
        fn = f"{f.replace('.bin', '')}.safetensors"

        if fn in files:
            print(f'Skipping, as {fn} already exists.')
            continue

        try:
            with torch.no_grad():
                state_dict = torch.load(f, map_location="cpu")
                id_encoder = state_dict["id_encoder"]
                to_be_changed = []
                for key in id_encoder:
                    #print(key) 
                    if "layrnorm" in key:
                        print(key)
                        newkey = key.replace("layrnorm", "layernorm")
                        to_be_changed.append((newkey, key))
                    if "visual_projection.w" in key:
                        print(key)
                        newkey = "vision_model."+key
                        to_be_changed.append((newkey, key))
                for nkey,okey in to_be_changed:
                    id_encoder[nkey] = id_encoder.pop(okey)
                lora = state_dict["lora_weights"]                
                weights =  id_encoder | lora
                #fn = f"{f.replace('.bin', '')}.safetensors"
                print(f'Saving {fn}...')
                save_file(weights, fn)
        except Exception as ex:
            print(f'ERROR converting {f}: {ex}')

print('Done!')