Repository: huggingface/diffusers

Issue:
train_text_to_image_sdxl.py ema not working
### 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_


Discussion:
I felt the same issue here.
Can you please add a reproducible code snippet? 
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)
```
hi @Jack000:
are you interested in open a PR? 
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](https://github.com/huggingface/diffusers/blob/main/CONTRIBUTING.md) are likely to be ignored.
So sorry for the delay here. Apologies. Could you please submit a PR fixing the issue? Looks like you have already found the bug. 
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](https://github.com/huggingface/diffusers/blob/main/CONTRIBUTING.md) are likely to be ignored.
Gentle bump :-) 

Resolving patch:
diff --git a/examples/text_to_image/train_text_to_image_sdxl.py b/examples/text_to_image/train_text_to_image_sdxl.py
--- a/examples/text_to_image/train_text_to_image_sdxl.py
+++ b/examples/text_to_image/train_text_to_image_sdxl.py
@@ -951,6 +951,9 @@ def collate_fn(examples):
         unet, optimizer, train_dataloader, lr_scheduler
     )
 
+    if args.use_ema:
+        ema_unet.to(accelerator.device)
+
     # We need to recalculate our total training steps as the size of the training dataloader may have changed.
     num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
     if overrode_max_train_steps:
@@ -1126,6 +1129,8 @@ def compute_time_ids(original_size, crops_coords_top_left):
 
             # Checks if the accelerator has performed an optimization step behind the scenes
             if accelerator.sync_gradients:
+                if args.use_ema:
+                    ema_unet.step(unet.parameters())
                 progress_bar.update(1)
                 global_step += 1
                 accelerator.log({"train_loss": train_loss}, step=global_step)


Test patch: