RuntimeError: Error(s) in loading state_dict for LatentDiffusion: size mismatch

[patryk-bartkowiak-nitid]

Hey, after I fine-tuned CyberRealistic model using DreamBooth I can't load it in AUTOMATIC1111

Error

Calculating sha256 for /project/stable-diffusion-webui/models/Stable-diffusion/del-250.safetensors: a813b3d75a85ea6a95a4823de4d1f0e653f10af8da6a38561a3ebe5cd95ad7c7
Loading weights [a813b3d75a] from /project/stable-diffusion-webui/models/Stable-diffusion/del-250.safetensors
Creating model from config: /project/stable-diffusion-webui/configs/v1-inference.yaml
changing setting sd_model_checkpoint to del-250.safetensors: RuntimeError
Traceback (most recent call last):
  File "/project/stable-diffusion-webui/modules/options.py", line 146, in set
    option.onchange()
  File "/project/stable-diffusion-webui/modules/call_queue.py", line 13, in f
    res = func(*args, **kwargs)
  File "/project/stable-diffusion-webui/modules/initialize_util.py", line 174, in <lambda>
    shared.opts.onchange("sd_model_checkpoint", wrap_queued_call(lambda: sd_models.reload_model_weights()), call=False)
  File "/project/stable-diffusion-webui/modules/sd_models.py", line 783, in reload_model_weights
    load_model(checkpoint_info, already_loaded_state_dict=state_dict)
  File "/project/stable-diffusion-webui/modules/sd_models.py", line 658, in load_model
    load_model_weights(sd_model, checkpoint_info, state_dict, timer)
  File "/project/stable-diffusion-webui/modules/sd_models.py", line 375, in load_model_weights
    model.load_state_dict(state_dict, strict=False)
  File "/project/stable-diffusion-webui/modules/sd_disable_initialization.py", line 223, in <lambda>
    module_load_state_dict = self.replace(torch.nn.Module, 'load_state_dict', lambda *args, **kwargs: load_state_dict(module_load_state_dict, *args, **kwargs))
  File "/project/stable-diffusion-webui/modules/sd_disable_initialization.py", line 221, in load_state_dict
    original(module, state_dict, strict=strict)
  File "/project/stable-diffusion-webui/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 LatentDiffusion:
        size mismatch for first_stage_model.encoder.mid.attn_1.q.bias: copying a param with shape torch.Size([512, 1, 1]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for first_stage_model.encoder.mid.attn_1.k.bias: copying a param with shape torch.Size([512, 1, 1]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for first_stage_model.encoder.mid.attn_1.v.bias: copying a param with shape torch.Size([512, 1, 1]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for first_stage_model.encoder.mid.attn_1.proj_out.bias: copying a param with shape torch.Size([512, 1, 1]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for first_stage_model.decoder.mid.attn_1.q.bias: copying a param with shape torch.Size([512, 1, 1]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for first_stage_model.decoder.mid.attn_1.k.bias: copying a param with shape torch.Size([512, 1, 1]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for first_stage_model.decoder.mid.attn_1.v.bias: copying a param with shape torch.Size([512, 1, 1]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for first_stage_model.decoder.mid.attn_1.proj_out.bias: copying a param with shape torch.Size([512, 1, 1]) from checkpoint, the shape in current model is torch.Size([512]).

Code used to fine-tune model

export MODEL_NAME="pretrained_models/cyberrealistic_v41"
export INSTANCE_DIR="data/cut_del"
export MAX_TRAIN_STEPS=750
export CHECKPOINTING_STEPS=125


export CLASS_DIR="${INSTANCE_DIR}_$(basename ${MODEL_NAME})_negative"
export OUTPUT_DIR="outputs/$(basename ${MODEL_NAME})_$(basename ${INSTANCE_DIR})"
export CUDA_LAUNCH_BLOCKING=1
export TORCH_USE_CUDA_DSA=1

printf "\n\nTraining with $MODEL_NAME on $INSTANCE_DIR, saving to $OUTPUT_DIR\n\n"

accelerate launch diffusers/examples/dreambooth/train_dreambooth.py \
  --instance_prompt="a photo of sts man, isolated on a white background" \
  --class_prompt="photo of a man, isolated on white background" \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --train_text_encoder \
  --instance_data_dir=$INSTANCE_DIR \
  --class_data_dir=$CLASS_DIR \
  --output_dir=$OUTPUT_DIR \
  --with_prior_preservation --prior_loss_weight=1.0 \
  --resolution=512 \
  --train_batch_size=4 \
  --use_8bit_adam \
  --gradient_checkpointing \
  --learning_rate=2e-6 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --num_class_images=200 \
  --max_train_steps=$MAX_TRAIN_STEPS \
  --checkpointing_steps=$CHECKPOINTING_STEPS \
  --enable_xformers_memory_efficient_attention \
  --seed="0"

Code used to convert checkpoint into suitable format:

import os
import pprint
import argparse

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", type=str, required=True)
    parser.add_argument("--name", type=str, required=True)
    parser.add_argument("--output", type=str, default="stable-diffusion-webui/models/Stable-diffusion/")
    args = parser.parse_args()

    print("#" * 100)
    print(f"INPUT: {args.input}")
    print(f"OUTPUT: {args.output}")
    print("#" * 100)
    print("\n" * 2)

    all_dirs = os.listdir(args.input)
    ckpt_dirs = sorted([d for d in all_dirs if d.startswith("checkpoint-") and os.path.isdir(os.path.join(args.input, d))])
    print("Found checkpoints:")
    pprint.pprint(ckpt_dirs)
    print("\n" * 2)

    for ckpt_dir in ckpt_dirs:
        input_path = os.path.abspath(os.path.join(args.input, ckpt_dir))
        output_path = os.path.abspath(
            os.path.join(args.output, args.name + '-' + ckpt_dir.split("-")[1] + '.safetensors')
        )

        print(f"Copying VAE model from {os.path.join(args.input, 'vae')} to {input_path}")
        os.system(f"cp -r {os.path.join(args.input, 'vae')} {input_path}")

        print(f"Converting to .safetensors format and copying to {output_path}")
        os.system(
            f"""python /project/diffusers/scripts/convert_diffusers_to_original_stable_diffusion.py \
                --use_safetensors \
                --model_path {input_path}\
                --checkpoint_path {output_path}"""
        )

If there are any additional info that I can provide to help detect the issue let me know!