# Got a bunch of .safetensors files to convert?
# Here's a handy script to take care of all that for you!
# Original .safetensors 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 .safetensors files, open a command prompt for that folder, and run:
# python convert_to_ckpt.py

import os
import torch
from safetensors.torch import load_file
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Using device:', device)
print()

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

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

        try:
            with torch.no_grad():
                weights = load_file(f)
                fn = f"{f.replace('.safetensors', '')}.ckpt"
                print(f'Saving {fn}...')
                torch.save(weights, fn)
        except Exception as ex:
            print(f'ERROR converting {f}: {ex}')

print('Done!')