Repository: huggingface/diffusers

Issue:
Possible misuse of scheduler in function log_validation in train_dreambooth.py
I copy train_dreambooth.py and use it to fully train IF Stage II. During training, I found it is unavoidable to generate images with noise point.I personally think it is due to the mandatory use of DPMSolverMultistepScheduler in function log_validation. When I try to use the origin scheduler DDPMScheduler offered by (https://huggingface.co/DeepFloyd/IF-II-M-v1.0/tree/main/scheduler), the problem is alleviated. Here are some comparisons using same checkpoint (above is DPMSolverMultistepScheduler).


![image](https://github.com/huggingface/diffusers/assets/43692893/4350b3d1-fd0e-40be-95bd-ee543cff93ee)
![image](https://github.com/huggingface/diffusers/assets/43692893/acca5036-9602-49b9-b1d2-c063b0a4bf70)

![image](https://github.com/huggingface/diffusers/assets/43692893/6cb61721-3db3-47b8-a2df-574bddbae712)
![image](https://github.com/huggingface/diffusers/assets/43692893/fb95e413-e8f4-448a-9729-de61cf5c8c3d)

Perhaps you could consider passing a new args parameter to control.


Discussion:
cc @sayakpaul and @williamberman for dreambooth 
Thanks so much for finding this!

Would you maybe like to open a PR for this? 
Yes, I similarly got better results while using the ddpm scheduler while training dreambooth IF.

Do we know if this is a bug in DPMSolverMultistepScheduler, if maybe it requires more steps for IF, or if the best solution is to just not diagnose and go with ddpm?
Thanks for your reply. I tried training twice and found this problem both times. One trained for a day, the other for two hours, both on one 8xA100. I think the DPMSolverMultistepScheduler is overly inclined to speed over quality, which may not be applied to an architecture without vae such as IF. I will open a PR for this.

Resolving patch:
diff --git a/examples/dreambooth/train_dreambooth.py b/examples/dreambooth/train_dreambooth.py
--- a/examples/dreambooth/train_dreambooth.py
+++ b/examples/dreambooth/train_dreambooth.py
@@ -17,6 +17,7 @@
 import copy
 import gc
 import hashlib
+import importlib
 import itertools
 import logging
 import math
@@ -47,7 +48,6 @@
     AutoencoderKL,
     DDPMScheduler,
     DiffusionPipeline,
-    DPMSolverMultistepScheduler,
     StableDiffusionPipeline,
     UNet2DConditionModel,
 )
@@ -144,7 +144,9 @@ def log_validation(
 
         scheduler_args["variance_type"] = variance_type
 
-    pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config, **scheduler_args)
+    module = importlib.import_module("diffusers")
+    scheduler_class = getattr(module, args.validation_scheduler)
+    pipeline.scheduler = scheduler_class.from_config(pipeline.scheduler.config, **scheduler_args)
     pipeline = pipeline.to(accelerator.device)
     pipeline.set_progress_bar_config(disable=True)
 
@@ -547,6 +549,13 @@ def parse_args(input_args=None):
         default=None,
         help="The optional `class_label` conditioning to pass to the unet, available values are `timesteps`.",
     )
+    parser.add_argument(
+        "--validation_scheduler",
+        type=str,
+        default="DPMSolverMultistepScheduler",
+        choices=["DPMSolverMultistepScheduler", "DDPMScheduler"],
+        help="Select which scheduler to use for validation. DDPMScheduler is recommended for DeepFloyd IF.",
+    )
 
     if input_args is not None:
         args = parser.parse_args(input_args)


Test patch: