Degraded Image Generation Performance in SD-XL U-Net Fine-tuning

[wddwzwhhxx]

Describe the bug

Problem:
I am experiencing degraded image generation performance while fine-tuning the U-Net part of the SDXL model using the train_text_to_image_sdxl.py script from the diffusers library. The generated results are noticeably worse compared to training on SD1.5 with the exact same dataset.
Here are some examples of generating indoor decoration images：
sd1.5:


sd-xl1.0:

Then I tried to finetune the sdxl on the pokeman dataset. As the number of training steps increased, the image generation effect seemed to deteriorate, and the fine-tuning script of the sdxl model seemed to only damage the original model's generation effect:
steps 100:

steps 300:

steps 1000:

steps 2000:

Has anyone else encountered a similar issue?
Alternatively, has anyone achieved favorable results when fine-tuning the U-Net script in SDXL?
I hope to receive insights from the community on this matter. Thank you!

Reproduction

export MODEL_NAME="stable-diffusion-xl-base-1.0"
export VAE_NAME="sdxl-vae-fp16-fix"
export DATASET_NAME="pokemon-blip-captions"

/aistudio/workspace/system-default/envs/diffusers/bin/python train_sdxl_test.py
--pretrained_model_name_or_path=$MODEL_NAME
--pretrained_vae_model_name_or_path=$VAE_NAME
--dataset_name=$DATASET_NAME
--output_dir=/aistudio/workspace/sd/diffusers/examples/text_to_image/ft_model/pokemon_modify/
--resolution=512 --center_crop --random_flip
--proportion_empty_prompts=0.2
--train_batch_size=1
--enable_xformers_memory_efficient_attention
--gradient_accumulation_steps=4 --gradient_checkpointing
--max_train_steps=2000
--learning_rate=1e-06 --lr_scheduler="constant" --lr_warmup_steps=0
--mixed_precision="fp16"
--use_8bit_adam
--checkpointing_steps=100
--dataloader_num_workers=20

Logs

No response

System Info

• diffusers version: 0.22.0.dev0
• Platform: Linux-4.19.96-x86_64-with-glibc2.17
• Python version: 3.10.13
• PyTorch version (GPU?): 2.1.0+cu121 (True)
• Huggingface_hub version: 0.17.3
• Transformers version: 4.35.1
• Accelerate version: 0.24.1
• xFormers version: 0.0.22.post7
• Using GPU in script?:
• Using distributed or parallel set-up in script?: yes

Who can help?

@sayakpaul @patrickvonplaten

[jorgemcgomes]

I have similar issues. I haven't been able to do full fine-tuning of SDXL without significantly degrading the quality of the output. Things start to get all messed up similar to what happened to you.

For what it's worth, I've tried using up to global batch size of 128 with LR=1e-6 and the problem remains.
And I've tried to fine-tune other SDXL variants/fine-tunes (including SSD-1B), and the same happens.

[sayakpaul]

I think it might be because of how the underlying UNet was pre-trained for SDXL. It tends to perform very poorly when we try to naively upsample smaller images to become 1024x0124, and it doesn't perform well on lower-resolution images either.

Then there are other training tricks (such as Min-SNR) that might help with the issue as well. I do think it's more of a hyperparameter tuning problem than a library issue.

So, I recommend fine-tuning it using LoRA.

[wddwzwhhxx]

I think it might be because of how the underlying UNet was pre-trained for SDXL. It tends to perform very poorly when we try to naively upsample smaller images to become 1024x0124, and it doesn't perform well on lower-resolution images either.

Then there are other training tricks (such as Min-SNR) that might help with the issue as well. I do think it's more of a hyperparameter tuning problem than a library issue.

So, I recommend fine-tuning it using LoRA.

Thank you for your response! Due to certain reasons, I need to fine-tune my model on a dataset with millions of samples. Could you guide me on adjusting my hyperparameters? I am open to generating lower-quality images, but how should I modify my script accordingly? Additionally, I am curious if LoRA is suitable for a dataset of this scale; I suspect there might be a performance limit. Thank you for your response!

[jorgemcgomes]

I think it might be because of how the underlying UNet was pre-trained for SDXL. It tends to perform very poorly when we try to naively upsample smaller images to become 1024x0124, and it doesn't perform well on lower-resolution images either.

Then there are other training tricks (such as Min-SNR) that might help with the issue as well. I do think it's more of a hyperparameter tuning problem than a library issue.

So, I recommend fine-tuning it using LoRA.

I'm not blaming the script, it might just be too sensitive to hyperparameters and there isn't much info around on that.

The low-resolution issue was my first guess as well. So I tried a SDXL model fine-tuned for 512x512 images: https://huggingface.co/hotshotco/SDXL-512
It does get you nice looking small images right from the beginning of the training process, but eventually the quality starts to degrade as well.

PS: while looking at the SDXL script today https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_sdxl.py I noticed the EMA implementation is a bit of a mess. The EMA weights are never updated and there is a missing unet weight restore after validation. So definitely don't use that @wddwzwhhxx

[sayakpaul]

PS: while looking at the SDXL script today https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_sdxl.py I noticed the EMA implementation is a bit of a mess. The EMA weights are never updated and there is a missing unet weight restore after validation. So definitely don't use that @wddwzwhhxx

Thanks for pointing that out! Happy to accept contributions for fixing that issue.

The low-resolution issue was my first guess as well. So I tried a SDXL model fine-tuned for 512x512 images: https://huggingface.co/hotshotco/SDXL-512
It does get you nice looking small images right from the beginning of the training process, but eventually the quality starts to degrade as well.

Yeah similar observations. Could make sense to do some kind of aspect-ratio bucketing to closely match how SDXL was originally trained. But our training scripts provide a simple basis, so we didn't incorporate that.

[jorgemcgomes]

Yeah similar observations. Could make sense to do some kind of aspect-ratio bucketing to closely match how SDXL was originally trained. But our training scripts provide a simple basis, so we didn't incorporate that.

I do have aspect-ratio bucketing in my training script (where I mix not only different aspect ratios but also high-res and low-res images). Doesn't help with this issue.
I've tried a lot of different things and haven't managed to find a combination that doesn't result in quality degradation after a few hundred training steps. My dataset is also quite large and diverse, similar to @wddwzwhhxx use case it seems.

And also similar to @wddwzwhhxx I have no issues training a SD1.5 or SD2.1 model with exactly the same dataset.

[jorgemcgomes]

Thanks for pointing that out! Happy to accept contributions for fixing that issue.

I already fixed it on my end and just wanted to do a training run to confirm it's working. I'll try to make a PR today/tomorrow.

[sayakpaul]

Thanks for sharing. Did you by chance use the faulty EMA updates?

For the low-res images, they are probably acting as outliers here. What happens when you train from scratch? I'd assume you have quite a lot of samples that could use for a from-scratch training run?

Another intuition would be because of how we're formulating the training objective. Maybe closely matching it would be more helpful. This is what I mean:

diffusers/examples/research_projects/controlnet/train_controlnet_webdataset.py

Line 1369 in 50a749e

model_pred = model_pred * (-sigmas) + noisy_latents

[wddwzwhhxx]

I think it might be because of how the underlying UNet was pre-trained for SDXL. It tends to perform very poorly when we try to naively upsample smaller images to become 1024x0124, and it doesn't perform well on lower-resolution images either.
Then there are other training tricks (such as Min-SNR) that might help with the issue as well. I do think it's more of a hyperparameter tuning problem than a library issue.
So, I recommend fine-tuning it using LoRA.

I'm not blaming the script, it might just be too sensitive to hyperparameters and there isn't much info around on that.

The low-resolution issue was my first guess as well. So I tried a SDXL model fine-tuned for 512x512 images: https://huggingface.co/hotshotco/SDXL-512 It does get you nice looking small images right from the beginning of the training process, but eventually the quality starts to degrade as well.

PS: while looking at the SDXL script today https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_sdxl.py I noticed the EMA implementation is a bit of a mess. The EMA weights are never updated and there is a missing unet weight restore after validation. So definitely don't use that @wddwzwhhxx

Thank you for sharing your experience on this issue! I confirm that I did not use EMA. The resolution might indeed be sensitive to the experimental results, But I haven't figured out how to modify the code yet.

[wddwzwhhxx]

closely matching

"I haven't used EMA. The cost of training from scratch seems too high, so I probably won't go that route... Regarding the 'closely matching' you mentioned, I will try it later on pokeman and other datasets. Thank you!"

[linnanwang]

I'm experiencing the similar issue, and watching how things goes here.

[wddwzwhhxx]

After multiple experiments with resolution issues, it seems I've come to some conclusions: The default training resolution in train_text_to_image_sdxl.py is 512. Therefore, when I set the resolution to 512 during inference, the generated results become normal (previously, I was inferring at the default resolution of 1024).


Additionally, I am also attempting to directly fine-tune on images with a resolution of 1024. I will update the results of this experiment later. @sayakpaul @jorgemcgomes