Repository: huggingface/diffusers

Issue:
Textual Inversion script does not resume from checkpoint
### Describe the bug

The textual inversion script `textual_inversion.py` located in `examples/textual_inversion` does not seem to resume properly from a checkpoint. If I interrupt the script halfway through training and then try to resume with the `--resume_from_checkpoint` argument, the script runs to completion without starting any training sets. 

I added additional logging and it seems part of the issue is that the `first_epoch` that is calculated during the resume process is greater than `num_epochs`, as shown in the log below.

### Reproduction

```
python textual_inversion.py \
       --pretrained_model_name_or_path=runwayml/stable-diffusion-v1-5 \
       --resolution=512 \
       --lr_scheduler=constant \
       --mixed_precision=fp16 \
       --learnable_property=object \
       --initializer_token='*' \
       --placeholder_token='<jello>' \
       --train_data_dir=/home/lstein/invokeai/training-data/jello \
       --output_dir=/home/lstein/invokeai/text-inversion-training/jello \
       --scale_lr \
       --train_batch_size=8 \
       --gradient_accumulation_steps=4 \
       --num_train_epochs=100 \
       --max_train_steps=3000 \
       --lr_warmup_steps=0 \
       --learning_rate=0.0005 \
       --resume_from_checkpoint=latest \
       --only_save_embeds
```

### Logs

```shell
01/20/2023 12:41:02 - INFO - ldm.invoke.textual_inversion_training - Distributed environment: NO
Num processes: 1
Process index: 0
Local process index: 0
Device: cuda
Mixed precision type: fp16

{'variance_type', 'prediction_type'} was not found in config. Values will be initialized to default values.
{'only_cross_attention', 'class_embed_type', 'use_linear_projection', 'upcast_attention', 'num_class_embeds', 'mid_block_type', 'dual_cross_attention', 'resnet_time_scale_shift'} was not found in config. Values will be initialized to default values.
01/20/2023 12:41:08 - INFO - ldm.invoke.textual_inversion_training - ***** Running training *****
01/20/2023 12:41:08 - INFO - ldm.invoke.textual_inversion_training -   Num examples = 600
01/20/2023 12:41:08 - INFO - ldm.invoke.textual_inversion_training -   Num Epochs = 158
01/20/2023 12:41:08 - INFO - ldm.invoke.textual_inversion_training -   Instantaneous batch size per device = 8
01/20/2023 12:41:08 - INFO - ldm.invoke.textual_inversion_training -   Total train batch size (w. parallel, distributed & accumulation) = 32
01/20/2023 12:41:08 - INFO - ldm.invoke.textual_inversion_training -   Gradient Accumulation steps = 4
01/20/2023 12:41:08 - INFO - ldm.invoke.textual_inversion_training -   Total optimization steps = 3000
Resuming from checkpoint checkpoint-2500
01/20/2023 12:41:08 - INFO - accelerate.accelerator - Loading states from /home/lstein/invokeai/text-inversion-training/jello/checkpoint-2500
01/20/2023 12:41:08 - INFO - accelerate.checkpointing - All model weights loaded successfully
01/20/2023 12:41:08 - INFO - accelerate.checkpointing - All optimizer states loaded successfully
01/20/2023 12:41:08 - INFO - accelerate.checkpointing - All scheduler states loaded successfully
01/20/2023 12:41:08 - INFO - accelerate.checkpointing - GradScaler state loaded successfully
01/20/2023 12:41:08 - INFO - accelerate.checkpointing - All random states loaded successfully
01/20/2023 12:41:08 - INFO - accelerate.accelerator - Loading in 0 custom states
01/20/2023 12:41:08 - INFO - ldm.invoke.textual_inversion_training -   Resume: resume_global_step = 10000
01/20/2023 12:41:08 - INFO - ldm.invoke.textual_inversion_training -   Resume: resume_step = 6
01/20/2023 12:41:08 - INFO - ldm.invoke.textual_inversion_training -   Resume: first_epoch = 526
Steps:   0%|                                                                                                                                                             | 0/500 [00:00<?, ?it/s]01/20/2023 12:41:08 - INFO - ldm.invoke.textual_inversion_training - Saving embeddings
Steps:   0%|                                                                                                                                                             | 0/500 [00:00<?, ?it/s]
>> Training completed. Copying learned_embeds.bin into /home/lstein/invokeai/embeddings/jello
```


### System Info

- `diffusers` version: 0.11.1
- Platform: Linux-5.15.0-56-generic-x86_64-with-glibc2.35
- Python version: 3.10.6
- PyTorch version (GPU?): 1.13.1+cu117 (True)
- Huggingface_hub version: 0.11.1
- Transformers version: 4.25.1
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?:  no


Discussion:
Thanks a lot for the issue. @pcuenca could you take a look here? 

Resolving patch:
diff --git a/examples/textual_inversion/textual_inversion.py b/examples/textual_inversion/textual_inversion.py
--- a/examples/textual_inversion/textual_inversion.py
+++ b/examples/textual_inversion/textual_inversion.py
@@ -597,7 +597,7 @@ def main():
         text_encoder, optimizer, train_dataloader, lr_scheduler
     )
 
-    # For mixed precision training we cast the text_encoder and vae weights to half-precision
+    # For mixed precision training we cast the unet and vae weights to half-precision
     # as these models are only used for inference, keeping weights in full precision is not required.
     weight_dtype = torch.float32
     if accelerator.mixed_precision == "fp16":
@@ -643,14 +643,21 @@ def main():
             dirs = os.listdir(args.output_dir)
             dirs = [d for d in dirs if d.startswith("checkpoint")]
             dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
-            path = dirs[-1]
-        accelerator.print(f"Resuming from checkpoint {path}")
-        accelerator.load_state(os.path.join(args.output_dir, path))
-        global_step = int(path.split("-")[1])
-
-        resume_global_step = global_step * args.gradient_accumulation_steps
-        first_epoch = resume_global_step // num_update_steps_per_epoch
-        resume_step = resume_global_step % num_update_steps_per_epoch
+            path = dirs[-1] if len(dirs) > 0 else None
+
+        if path is None:
+            accelerator.print(
+                f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
+            )
+            args.resume_from_checkpoint = None
+        else:
+            accelerator.print(f"Resuming from checkpoint {path}")
+            accelerator.load_state(os.path.join(args.output_dir, path))
+            global_step = int(path.split("-")[1])
+
+            resume_global_step = global_step * args.gradient_accumulation_steps
+            first_epoch = global_step // num_update_steps_per_epoch
+            resume_step = resume_global_step % (num_update_steps_per_epoch * args.gradient_accumulation_steps)
 
     # Only show the progress bar once on each machine.
     progress_bar = tqdm(range(global_step, args.max_train_steps), disable=not accelerator.is_local_main_process)


Test patch: