wip

Author: williamberman

examples/dreambooth/train_dreambooth_lora.py

@@ -44,7 +44,6 @@
     AutoencoderKL,
     DDPMScheduler,
     DiffusionPipeline,
-    DPMSolverMultistepScheduler,
     StableDiffusionPipeline,
     UNet2DConditionModel,
 )
@@ -1092,7 +1091,11 @@ def compute_text_embeddings(prompt):
                     revision=args.revision,
                     torch_dtype=weight_dtype,
                 )
-                pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
+
+                # TODO temp hack for IF
+                # pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
+                pipeline.scheduler = DDPMScheduler.from_config(pipeline.scheduler.config, variance_type="fixed_small")
+
                 pipeline = pipeline.to(accelerator.device)
                 pipeline.set_progress_bar_config(disable=True)
 
@@ -1143,7 +1146,11 @@ def compute_text_embeddings(prompt):
         pipeline = DiffusionPipeline.from_pretrained(
             args.pretrained_model_name_or_path, revision=args.revision, torch_dtype=weight_dtype
         )
-        pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
+
+        # TODO temp for IF
+        # pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
+        pipeline.scheduler = DDPMScheduler.from_config(pipeline.scheduler.config, variance_type="fixed_small")
+
         pipeline = pipeline.to(accelerator.device)
 
         # load attention processors
@@ -1153,7 +1160,7 @@ def compute_text_embeddings(prompt):
         if args.validation_prompt and args.num_validation_images > 0:
             generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed else None
             images = [
-                pipeline(args.validation_prompt, num_inference_steps=25, generator=generator).images[0]
+                pipeline(args.validation_prompt, generator=generator).images[0]
                 for _ in range(args.num_validation_images)
             ]