Finetuned models aren't saved or loaded properly in train_custom_diffusion.py

[mostafij-rahman]

Describe the bug

Thank you for your amazing work. It seems like models are not saved or loaded properly after finetuning train_custom_diffusion.py in a new dataset. Generated validation images are expected but final test (inference) results are not expected as finetuned models are not used properly. The bug most probably in the following lines.

Reproduction

Please finetune using a custom dataset and see the final test/inference results.

Logs

No response

System Info

• 珞 Diffusers version: 0.31.0.dev0
• Platform: Linux-6.5.0-45-generic-x86_64-with-glibc2.17
• Running on Google Colab?: No
• Python version: 3.8.19
• PyTorch version (GPU?): 2.4.0+cu121 (True)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Huggingface_hub version: 0.24.6
• Transformers version: 4.44.0
• Accelerate version: 0.33.0
• PEFT version: not installed
• Bitsandbytes version: not installed
• Safetensors version: 0.4.4
• xFormers version: not installed
• Accelerator: NVIDIA RTX A6000, 49140 MiB
  NVIDIA RTX A6000, 49140 MiB
  NVIDIA RTX A6000, 49140 MiB
  NVIDIA RTX A6000, 49140 MiB
• Using GPU in script?: 4 NVIDIA RTX A6000 GPUs
• Using distributed or parallel set-up in script?: distributed

Who can help?

No response

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[E774-ChangxingYang37]

Hi, have you solved this issue?

[liulisixin]

Hi, thank you for your issue. I also met this problem.

Does anyone solve this problem or is it fixed?

[liulisixin]

After checking for some times, I found the following codes from diffusers may contain bugs:

unet.save_attn_procs(args.output_dir, safe_serialization=not args.no_safe_serialization)
pipeline.unet.load_attn_procs(args.output_dir, weight_name=weight_name)

After saving and loading from the files, the processors for CustomDiffusion are all lost.

Now, I avoided it by downgrading the versions:

pip install diffusers==0.26
pip install huggingface-hub==0.20.2
pip install transformers==4.28

If I have time in the future, I want to look into the unet.save_attn_procs and unet.load_attn_procs to see what is wrong.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[JStyborski]

I think I figured out the bug.

In later versions of diffusers (I'm using 0.32.2) the pipeline.unet.load_attn_procs( ) script adds an if loop to the CustomDiffusion process (see attached image). Now you must additionally pass the _pipeline argument. Without passing _pipeline argument, load_attn_procs( ) will load the finetuned CustomDiffusion weights, but it will never set them. Since the attn_processors are neither set nor returned, they are lost when when load_attn_procs( ) returns. Adding the _pipeline argument will trigger the loop that calls set_attn_processor( ).

However, I think there's an additional bug here. At least in my case, CustomDiffusion saves with 16 procs (corresponding to each of 16 cross-attention layers). However, models.unets.unet_2d_condition.set_attn_processor( ) has a check to ensure that the number of procs matches the number of attention layers in pipeline.unet.attn_processors, which contains 32 procs (corresponding to 16 cross-attention procs and 16 self-attention procs). This causes the set_attn_processor( ) call to fail in my case.

I created a workaround as follows:

• Load CustomDiffusion attention processors (16 procs)
• Add in self-attention processors from original model (to get to 32 procs)
• Call set_attn_processor on the mixed processor dictionary

This code allows me to successfully generate images with my finetuned CustomDiffusion model.

Final comment: I think HuggingFace is not really paying attention to CustomDiffusion anymore. This CustomDiffusion code has not been updated across many versions, yet this thread is the only one I see about the saving/loading bug. I feel like if someone were actively monitoring this, the bug would have been noticed/fixed sooner. Perhaps it is time to start refocusing on more popular finetuning methods like LoRAs.

[tuananhbui89]

I tried the solution from @JStyborski but didn't work from my side (error at line pipeline.unet._process_custom_diffusion(state_dict=custom_unet_state_dict), 'to_k_custom_diffusion.weight' is not in the state_dict)

I think there is a bug in the function _process_custom_diffusion where only one type of processor CustomDiffusionAttnProcessor was checked while in the _get_custom_diffusion_state_dict function, there are three types of processors to be checked.