KeyError: ‘global_step’ in File “[...]convert_[sd]”, line 799, in global_step = checkpoint[“global_step”]

And I found some custom ckpt files converted to diffuser well with code
but some custom ckpt files got error
Traceback (most recent call last):
File “/api/./diffusers/scripts/”, line 799, in
global_step = checkpoint[“global_step”]
KeyError: ‘global_step’

Maybe need other technic to change these not changed ckpt files?

Also moving this to it’s own topic :grin:

You’re doing this via the CHECKPOINT_URL build-arg, right?
Or some other way?

I haven’t encountered this before… if you’re using the latest dev commit, which uses (almost) the latest diffusers commit, it might be worth reporting to diffusers at Issues · huggingface/diffusers · GitHub (after first searching for any similar issues). Make sure to give them a link to the checkpoint you’re trying to convert.

I found what’s the problem with huggingface diffuser’s conver to diffuser script.
Problem occurs with finetuned models. The solution is change script 857-879 lines to

checkpoint = torch.load(args.checkpoint_path)

# Sometimes models don't have the global_step item
if "global_step" in checkpoint:
    global_step = checkpoint["global_step"]
    print("global_step key not found in model")
    global_step = None

chckpoint_dict_replacements = {
'cond_stage_model.transformer.embeddings.': 'cond_stage_model.transformer.text_model.embeddings.',
'cond_stage_model.transformer.encoder.': 'cond_stage_model.transformer.text_model.encoder.',
'cond_stage_model.transformer.final_layer_norm.': 'cond_stage_model.transformer.text_model.final_layer_norm.',

def transform_checkpoint_dict_key(k):
    for text, replacement in chckpoint_dict_replacements.items():
        if k.startswith(text):
            k = replacement + k[len(text):]

    return k

def get_state_dict_from_checkpoint(pl_sd):
    pl_sd = pl_sd.pop("state_dict", pl_sd)
    pl_sd.pop("state_dict", None)

    sd = {}
    for k, v in pl_sd.items():
        new_key = transform_checkpoint_dict_key(k)

        if new_key is not None:
            sd[new_key] = v


    return pl_sd

def read_state_dict(checkpoint_file):
    sd = get_state_dict_from_checkpoint(checkpoint_file)
    return sd

checkpoint = read_state_dict(checkpoint)

this leads to read state_dict correctly with any checkpoint. FP16 checkpoints turns to FP32 when finished, so if needs to do half() to unet, vae, text_encoder to do so. Find the way with AUTOMATIC1111 repo.

1 Like

Oh amazing! Thanks so much for sharing this for everyone!

I thiiiink I may have seen something about this too in diffuser’s latest commits… but I want to finish all the new testing infrastructure before I upgrade diffusers and make a new release.

In any event, great job!! :pray: :+1:
And most importantly, thanks for sharing.

1 Like