When using .ckpt, "comfy/sd.py": KeyError: 'state_dict'

[valconius]

This will probably be a long-term issue others encounter and can serve as a convenient reference. It might not be a ComfyUI issue at all. tl;dr Running the error'ing model through any model merger or converter (such as WebUI's) fixes the issue, see below for details.

When using some models, but not others, seemingly only model merges created months ago in WebUI, the following error occurs:

making attention of type 'vanilla' with 512 in_channels
Working with z of shape (1, 4, 32, 32) = 4096 dimensions.
making attention of type 'vanilla' with 512 in_channels
Loading model from /home/user/ComfyUI/models/checkpoints/model1.ckpt
Traceback (most recent call last):
  File "/home/user/ComfyUI/main.py", line 160, in execute
    executed += recursive_execute(prompt, self.outputs, x, extra_data)
  File "/home/user/ComfyUI/main.py", line 54, in recursive_execute
    executed += recursive_execute(prompt, outputs, input_unique_id, extra_data)
  File "/home/user/ComfyUI/main.py", line 54, in recursive_execute
    executed += recursive_execute(prompt, outputs, input_unique_id, extra_data)
  File "/home/user/ComfyUI/main.py", line 54, in recursive_execute
    executed += recursive_execute(prompt, outputs, input_unique_id, extra_data)
  File "/home/user/ComfyUI/main.py", line 59, in recursive_execute
    outputs[unique_id] = getattr(obj, obj.FUNCTION)(**input_data_all)
  File "/home/user/ComfyUI/nodes.py", line 83, in load_checkpoint
    return comfy.sd.load_checkpoint(config_path, ckpt_path, output_vae=True, output_clip=True)
  File "/home/user/ComfyUI/comfy/sd.py", line 123, in load_checkpoint
    model = load_model_from_config(config, ckpt_path, verbose=False, load_state_dict_to=load_state_dict_to)
  File "/home/user/ComfyUI/comfy/sd.py", line 20, in load_model_from_config
    sd = pl_sd["state_dict"]
KeyError: 'state_dict'

A reminder, the code for pl_sd is...

def load_model_from_config ...
    ...
    if ckpt.lower().endswith(".safetensors"):
        import safetensors.torch
        sd = safetensors.torch.load_file(ckpt, device="cpu")
    else:
        pl_sd = torch.load(ckpt, map_location="cpu")
        if "global_step" in pl_sd:
            print(f"Global Step: {pl_sd['global_step']}")
        sd = pl_sd["state_dict"] # <-- PROBLEM OCCCURS HERE
    model = instantiate_from_config(config.model)
...
model = load_model_from_config(config, ckpt_path, verbose=False, load_state_dict_to=load_state_dict_to)

... and as the error says, ckpt and ckpt_path are certainly correct.

This error is resolved by entering WebUI and running either of the following, after which ComfyUI can use the model. FP16 and Pruning are not necessary, but just provided as examples (as they are now the default used in WebUI's model merge utility):

• Model Converter Extension -> fp16 + no-ema prune (although anything would probably work)
• Model Merge -> Weighted Sum, model1 + model1, weight = 0.0, fp16
• CLI Pruner & Converter: https://github.com/lopho/stable-diffusion-prune

It seems the offending models were made with PyTorch 12.x, and both WebUI and ComfyUI are running on PyTorch 13.1. WebUI can read these (malformed?) ckpt files, but ComfyUI cannot.

EDIT: This seems to be a common problem with a simple compatibility fix:

• lopho/stable-diffusion-prune@38e3c10

[comfyanonymous]

acdc6f4

Should be fixed now.

[valconius]

Well holy cow that was a fast fix. Indeed, my broken models work again. I wonder:

• How could a provisional fix be put on git so that rather than doing "git pull" a branch could be pulled instead?
• Is this model breakage even okay, or is WebUI incorrectly saving ckpt files during model merge? I noticed even my newer merges have this problem; only merges and block merges from the following lack this problem:
  • https://github.com/bbc-mc/sdweb-merge-block-weighted-gui
  • https://github.com/bbc-mc/sdweb-merge-board

[comfyanonymous]

How could a provisional fix be put on git so that rather than doing "git pull" a branch could be pulled instead?

You can switch to another branch then: git pull https://github.com/comfyanonymous/ComfyUI other_branch
Replacing the repo url or branch name with the right one.

I don't think this model format is ok because it makes them incompatible with upstream SD:
https://github.com/Stability-AI/stablediffusion/blob/main/scripts/txt2img.py#L28