Convert all CKPT files to SAFETENSOR files in a directory, NAI fix

convert_to_safe_fix.py

# 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
#
# Original script by @xrpgame https://github.com/xrpgame
# https://gist.github.com/xrpgame/8f756f99b00b02697edcd5eec5202c59
#
# Edited by @RassilonSleeps for NAI model compatibility https://github.com/RassilonSleeps
# https://gist.github.com/RassilonSleeps/edfb630819b95307270efa8450163bc1

import os
import torch
from safetensors.torch import save_file

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

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

        try:
            with torch.no_grad():
                weights = torch.load(f)["state_dict"]
                weights.pop("state_dict")
                fn = f"{f.replace('.ckpt', '')}.safetensors"
                print(f'Saving {fn}...')
                save_file(weights, fn)
        except Exception as ex:
            print(f'ERROR converting {f}: {ex}')

print('Done!')

Your nai fix breaks all other conversions. Line 35 must be conditioned to only be run on NAI models. Fx by checking the file name or size, or the content of state_dict.

When the reception is thrown before saving the file, Python moves on top the next iteration in the For loop so you only get the NAI model saved, and nothing else... I posted a fix idea here: https://gist.github.com/xrpgame/8f756f99b00b02697edcd5eec5202c59?permalink_comment_id=4404713#gistcomment-4404713