train_text_to_image_sdxl.py ema not working

[Jack000]

Describe the bug

I did a few training runs with train_text_to_image_sdxl.py before realizing that the EMA checkpoint never changed.

Looking at train_text_to_image.py - it seems that there is no ema_unet.step() in the sdxl training code so the ema model is never updated.

it's also missing this bit at the start:

if args.use_ema:
    ema_unet.to(accelerator.device)

Reproduction

just the huggingface sample code

Logs

No response

System Info

latest version of diffusers

Who can help?

No response

[linnanwang]

I felt the same issue here.

[patrickvonplaten]

Can you please add a reproducible code snippet?

[Jack000]

I guess just the sample code, but with --use_ema turned on
from https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/README_sdxl.md

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

accelerate launch train_text_to_image_sdxl.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --pretrained_vae_model_name_or_path=$VAE_NAME \
  --dataset_name=$DATASET_NAME \
  --enable_xformers_memory_efficient_attention \
  --resolution=512 --center_crop --random_flip \
  --proportion_empty_prompts=0.2 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 --gradient_checkpointing \
  --max_train_steps=10000 \
  --use_8bit_adam \
  --learning_rate=1e-06 --lr_scheduler="constant" --lr_warmup_steps=0 \
  --mixed_precision="fp16" \
  --report_to="wandb" \
  --validation_prompt="a cute Sundar Pichai creature" --validation_epochs 5 \
  --checkpointing_steps=5000 \
  --output_dir="sdxl-pokemon-model" \
  --use_ema

this will error out with "RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!"

because it's missing

if args.use_ema:
    ema_unet.to(accelerator.device)

[yiyixuxu]

hi @Jack000:
are you interested in open a PR?

[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.

[sayakpaul]

So sorry for the delay here. Apologies. Could you please submit a PR fixing the issue? Looks like you have already found the bug.

[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.

[sayakpaul]

Gentle bump :-)

[tolgacangoz]

Since nobody has done that for a while; I tried.