[bug]: Cannot load legacy checkpoint, diffusers conversion fails

[bedalton]

Is there an existing issue for this?

• [X] I have searched the existing issues

OS

macOS

GPU

mps

VRAM

No response

What version did you experience this issue on?

2.3.5.post2 and 2.3.5.rc2(I think before I updated just now)

What happened?

When loading a new checkpoint trained with EveryDream v2, the model fails to load, and model selection reverts back to previous. Checkpoint loads fine in both Automatic1111 and Draw Things (macOS) and both work as expected.

Also, despite disabling legacy checkpoint conversion to diffusers in config, console still prints >> Converting legacy checkpoint dmemxt-v5 into a diffusers model...

The error output is here:

>> Model change requested: dmemxt-v5
>> Current VRAM usage:  0.00G
>> Cache limit (max=1) reached. Purging stable-diffusion-1.5
>> Loading dmemxt-v5 from /Users/..../models/StableDiffusion/dmemxt-v5.ckpt
>> Converting legacy checkpoint dmemxt-v5 into a diffusers model...
   | global_step key not found in model
   | Using checkpoint model's original VAE
** model dmemxt-v5 could not be loaded: Error(s) in loading state_dict for CLIPTextModel:
	Missing key(s) in state_dict: "text_model.embeddings.position_ids". 
Traceback (most recent call last):
  File "/Users/.../InvokeAI/.venv/lib/python3.10/site-packages/ldm/generate.py", line 964, in set_model
    model_data = cache.get_model(model_name)
  File "/Users/../InvokeAI/.venv/lib/python3.10/site-packages/ldm/invoke/model_manager.py", line 105, in get_model
    requested_model, width, height, hash = self._load_model(model_name)
  File "/Users/.../InvokeAI/.venv/lib/python3.10/site-packages/ldm/invoke/model_manager.py", line 317, in _load_model
    model, width, height, model_hash = self._load_ckpt_model(
  File "/Users/.../InvokeAI/.venv/lib/python3.10/site-packages/ldm/invoke/model_manager.py", line 376, in _load_ckpt_model
    pipeline = load_pipeline_from_original_stable_diffusion_ckpt(
  File "/Users/.../InvokeAI/.venv/lib/python3.10/site-packages/ldm/invoke/ckpt_to_diffuser.py", line 1296, in load_pipeline_from_original_stable_diffusion_ckpt
    text_model = convert_ldm_clip_checkpoint(checkpoint)
  File "/Users/.../InvokeAI/.venv/lib/python3.10/site-packages/ldm/invoke/ckpt_to_diffuser.py", line 860, in convert_ldm_clip_checkpoint
    text_model.load_state_dict(text_model_dict)
  File "/Users/.../InvokeAI/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 2041, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for CLIPTextModel:
	Missing key(s) in state_dict: "text_model.embeddings.position_ids".

Screenshots

No response

Additional context

Only happens with EveryDream2 models trained yesterday and today on RunPod, so could be something that changed in EveryDream. Older EveryDream checkpoints load fine in InvokeAI.

New checkpoints works in Automatic and Draw Things

Contact Details

No response

[bedalton]

Some checkpoints when loading, say they have missing keys like
encoder.mid_block.attentions.0.query.weight and have unexpected keys like:
encoder.mid_block.attentions.0.to_query.weight

So expecting query, but finding to_query

or encoder.mid_block.attentions.0.key.weight but having encoder.mid_block.attentions.0.to_key.weight

So key vs to_key

Error:

>> Model change requested: dmemxt-v5-ep150-gs01050
>> Current VRAM usage:  0.00G
>> Cache limit (max=1) reached. Purging c1-v10-1.4e-7-ep40-gs01480
>> Loading dmemxt-v5-ep150-gs01050 from /Users/../models/StableDiffusion/dmemxt-v5-ep150-gs01050.ckpt
>> Converting legacy checkpoint dmemxt-v5-ep150-gs01050 into a diffusers model...
   | Scanning Model: /Users/../models/StableDiffusion/dmemxt-v5-ep150-gs01050.ckpt
   | Model scanned ok
   | global_step key not found in model
   | Using checkpoint model's original VAE
** model dmemxt-v5-ep150-gs01050 could not be loaded: Error(s) in loading state_dict for AutoencoderKL:
	Missing key(s) in state_dict: "encoder.mid_block.attentions.0.query.weight", 
"encoder.mid_block.attentions.0.query.bias", "encoder.mid_block.attentions.0.key.weight", 
"encoder.mid_block.attentions.0.key.bias", "encoder.mid_block.attentions.0.value.weight", 
"encoder.mid_block.attentions.0.value.bias", "encoder.mid_block.attentions.0.proj_attn.weight", 
"encoder.mid_block.attentions.0.proj_attn.bias", "decoder.mid_block.attentions.0.query.weight", 
"decoder.mid_block.attentions.0.query.bias", "decoder.mid_block.attentions.0.key.weight", 
"decoder.mid_block.attentions.0.key.bias", "decoder.mid_block.attentions.0.value.weight", 
"decoder.mid_block.attentions.0.value.bias", "decoder.mid_block.attentions.0.proj_attn.weight", 
"decoder.mid_block.attentions.0.proj_attn.bias". 
	Unexpected key(s) in state_dict: "encoder.mid_block.attentions.0.to_query.weight", "encoder.mid_block.attentions.0.to_query.bias", "encoder.mid_block.attentions.0.to_key.weight", 
"encoder.mid_block.attentions.0.to_key.bias", "encoder.mid_block.attentions.0.to_value.weight", "encoder.mid_block.attentions.0.to_value.bias", "encoder.mid_block.attentions.0.to_out.0.weight", 
"encoder.mid_block.attentions.0.to_out.0.bias", "decoder.mid_block.attentions.0.to_query.weight", 
"decoder.mid_block.attentions.0.to_query.bias", "decoder.mid_block.attentions.0.to_key.weight", 
"decoder.mid_block.attentions.0.to_key.bias", "decoder.mid_block.attentions.0.to_value.weight", 
"decoder.mid_block.attentions.0.to_value.bias", "decoder.mid_block.attentions.0.to_out.0.weight", 
"decoder.mid_block.attentions.0.to_out.0.bias". 
Traceback (most recent call last):
  File "/Users/../InvokeAI/.venv/lib/python3.10/site-packages/ldm/generate.py", line 964, in set_model
    model_data = cache.get_model(model_name)
  File "/Users/../InvokeAI/.venv/lib/python3.10/site-packages/ldm/invoke/model_manager.py", line 105, in get_model
    requested_model, width, height, hash = self._load_model(model_name)
  File "/Users/../InvokeAI/.venv/lib/python3.10/site-packages/ldm/invoke/model_manager.py", line 317, in _load_model
    model, width, height, model_hash = self._load_ckpt_model(
  File "/Users/../InvokeAI/.venv/lib/python3.10/site-packages/ldm/invoke/model_manager.py", line 376, in _load_ckpt_model
    pipeline = load_pipeline_from_original_stable_diffusion_ckpt(
  File "/Users/../InvokeAI/.venv/lib/python3.10/site-packages/ldm/invoke/ckpt_to_diffuser.py", line 1253, in load_pipeline_from_original_stable_diffusion_ckpt
    vae.load_state_dict(converted_vae_checkpoint)
  File "/Users/../InvokeAI/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 2041, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for AutoencoderKL:
	Missing key(s) in state_dict: "encoder.mid_block.attentions.0.query.weight", 
"encoder.mid_block.attentions.0.query.bias", "encoder.mid_block.attentions.0.key.weight", 
"encoder.mid_block.attentions.0.key.bias", "encoder.mid_block.attentions.0.value.weight", 
"encoder.mid_block.attentions.0.value.bias", "encoder.mid_block.attentions.0.proj_attn.weight", 
"encoder.mid_block.attentions.0.proj_attn.bias", "decoder.mid_block.attentions.0.query.weight", 
"decoder.mid_block.attentions.0.query.bias", "decoder.mid_block.attentions.0.key.weight", 
"decoder.mid_block.attentions.0.key.bias", "decoder.mid_block.attentions.0.value.weight", 
"decoder.mid_block.attentions.0.value.bias", "decoder.mid_block.attentions.0.proj_attn.weight", 
"decoder.mid_block.attentions.0.proj_attn.bias". 
	Unexpected key(s) in state_dict: "encoder.mid_block.attentions.0.to_query.weight", 
"encoder.mid_block.attentions.0.to_query.bias", "encoder.mid_block.attentions.0.to_key.weight", 
"encoder.mid_block.attentions.0.to_key.bias", "encoder.mid_block.attentions.0.to_value.weight", 
"encoder.mid_block.attentions.0.to_value.bias", "encoder.mid_block.attentions.0.to_out.0.weight", 
"encoder.mid_block.attentions.0.to_out.0.bias", "decoder.mid_block.attentions.0.to_query.weight", 
"decoder.mid_block.attentions.0.to_query.bias", "decoder.mid_block.attentions.0.to_key.weight", 
"decoder.mid_block.attentions.0.to_key.bias", "decoder.mid_block.attentions.0.to_value.weight", 
"decoder.mid_block.attentions.0.to_value.bias", "decoder.mid_block.attentions.0.to_out.0.weight", 
"decoder.mid_block.attentions.0.to_out.0.bias". 

[n00mkrad]

@bedalton @hipsterusername How was this solved?

[hipsterusername]

@bedalton @hipsterusername How was this solved?

To be candid, we did a massive issue clean-up for anything pre-3.0. I would not be able to attest to this working without follow-up/re-testing, but I haven't seen many reports in discord, and we have a few Everydream contributors that are regulars.

[n00mkrad]

@bedalton @hipsterusername How was this solved?

To be candid, we did a massive issue clean-up for anything pre-3.0. I would not be able to attest to this working without follow-up/re-testing, but I haven't seen many reports in discord, and we have a few Everydream contributors that are regulars.

I managed to make a rough patch myself. In my case, I tried to use 2.3.5 with a newer Diffusers version which had changed some key names.

To fix it, I edited model_manager.py from this:

        vae = AutoencoderKL(**vae_config)
        vae.load_state_dict(diffusers_vae)
        return vae

to this:

        vae = AutoencoderKL(**vae_config)
        try:
            vae.load_state_dict(diffusers_vae)
        except:
            diffusers_vae = { key.replace(".query.", ".to_q.").replace(".key.", ".to_k.").replace(".value.", ".to_v.").replace(".proj_attn.", ".to_out.0."): value for key, value in diffusers_vae.items() }
            vae.load_state_dict(diffusers_vae)
        return vae

[hipsterusername]

Raising for visibility up to @lstein

[bedalton]

In my case the code posted by n00mkrad never gets called as I have is_accelerate_available, but it gave me an idea to change the attributes myself. In my case there are added parameters to_to_q, to_to_v, and to_to_k, which caused errors. So I modified convert_ckpt_to_diffusers and on line 178 after:

https://github.com/invoke-ai/InvokeAI/blob/ae17d01e1d2b4f2ef85a5df6c6e4d7ce0f378ca9/invokeai/backend/model_management/convert_ckpt_to_diffusers.py#L176-L177

I added:

if "to_to_q" in new_item:
    new_item = new_item.replace("to_to_q.weight", "to_q.weight")
    new_item = new_item.replace("to_to_q.bias", "to_q.bias")
if "to_to_k" in new_item:
    new_item = new_item.replace("to_to_k.weight", "to_k.weight")
    new_item = new_item.replace("to_to_k.bias", "to_k.bias")
if "to_to_v" in new_item:
    new_item = new_item.replace("to_to_v.weight", "to_v.weight")
    new_item = new_item.replace("to_to_v.bias", "to_v.bias")

This mostly works, but twice I had to adjust VAE and VAE precision to get it to work without problems

[bedalton]

Definitely an EveryDreamV2 problem, so I plan to let them know

[InB4DevOps]

Definitely an EveryDreamV2 problem, so I plan to let them know

I have tried it all and the only working solution I found is this: https://github.com/s9roll7/animatediff-cli-prompt-travel/issues/38#issuecomment-1732512147