UNet2DConditionModel keys missing error

[ihojun887]

Hi, thanks for your great work!

I ran into the following error when training from scratch with train_xl.py and then running inference with inference.py:

ValueError: Cannot load <class 'src.unet_hacked_tryon.UNet2DConditionModel'> because the following keys are missing:
up_blocks.0.attentions.2.transformer_blocks.4.attn2.processor.to_v_ip.weight, ...

Training command
CUDA_VISIBLE_DEVICES=0 accelerate launch train_xl.py
--gradient_checkpointing --use_8bit_adam
--output_dir=result2 --train_batch_size=1
--data_dir=~~~
--mixed_precision="bf16"
--enable_xformers_memory_efficient_attention

Inference command
PYTHONWARNINGS="ignore::FutureWarning" CUDA_VISIBLE_DEVICES=0 accelerate launch inference.py
--pretrained_model_name_or_path ~~~~
--width 768 --height 1024 --num_inference_steps 30
--output_dir IDM_result_paired
--data_dir ~~~
--seed 42
--test_batch_size 4
--guidance_scale 2.0

Question

Is this error caused by using --mixed_precision="bf16" and --enable_xformers_memory_efficient_attention during training?
Or is there something else I should check to make the model save and load properly?

Thanks in advance!