From 9199d698f8200783c577bbbd83acf6d14fef029a Mon Sep 17 00:00:00 2001
From: Kevin Turner <83819+keturn@users.noreply.github.com>
Date: Thu, 8 Dec 2022 13:02:47 -0800
Subject: [PATCH] diffusers: upgrade to diffusers 0.10, add Heun scheduler

---
 binary_installer/requirements.in           | 2 +-
 ldm/generate.py                            | 8 +++++++-
 ldm/invoke/generator/diffusers_pipeline.py | 4 ++--
 3 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/binary_installer/requirements.in b/binary_installer/requirements.in
index cada6ec51e..95ef3f5185 100644
--- a/binary_installer/requirements.in
+++ b/binary_installer/requirements.in
@@ -4,7 +4,7 @@
 --trusted-host https://download.pytorch.org
 accelerate~=0.14
 albumentations
-diffusers[torch]~=0.9
+diffusers[torch]~=0.10
 einops
 eventlet
 flask_cors
diff --git a/ldm/generate.py b/ldm/generate.py
index aefe58c8b7..82fe52f4ae 100644
--- a/ldm/generate.py
+++ b/ldm/generate.py
@@ -17,6 +17,7 @@ import skimage
 import torch
 import transformers
 from PIL import Image, ImageOps
+from diffusers import HeunDiscreteScheduler
 from diffusers.pipeline_utils import DiffusionPipeline
 from diffusers.schedulers.scheduling_ddim import DDIMScheduler
 from diffusers.schedulers.scheduling_dpmsolver_multistep import DPMSolverMultistepScheduler
@@ -1008,12 +1009,17 @@ class Generate:
 
         scheduler_map = dict(
             ddim=DDIMScheduler,
+            dpmpp_2=DPMSolverMultistepScheduler,
             ipndm=IPNDMScheduler,
+            # DPMSolverMultistepScheduler is technically not `k_` anything, as it is neither
+            # the k-diffusers implementation nor included in EDM (Karras 2022), but we can
+            # provide an alias for compatibility.
+            k_dpmpp_2=DPMSolverMultistepScheduler,
             k_euler=EulerDiscreteScheduler,
             k_euler_a=EulerAncestralDiscreteScheduler,
+            k_heun=HeunDiscreteScheduler,
             k_lms=LMSDiscreteScheduler,
             plms=PNDMScheduler,
-            k_dpmpp_2=DPMSolverMultistepScheduler,
         )
 
         if self.sampler_name in scheduler_map:
diff --git a/ldm/invoke/generator/diffusers_pipeline.py b/ldm/invoke/generator/diffusers_pipeline.py
index 4dbd576c51..8f2b4901b8 100644
--- a/ldm/invoke/generator/diffusers_pipeline.py
+++ b/ldm/invoke/generator/diffusers_pipeline.py
@@ -407,7 +407,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
         batch_size = initial_latents.size(0)
         img2img_pipeline = StableDiffusionImg2ImgPipeline(**self.components)
         img2img_pipeline.scheduler.set_timesteps(num_inference_steps, device=device)
-        timesteps = img2img_pipeline.get_timesteps(num_inference_steps, strength, device=device)
+        timesteps, _ = img2img_pipeline.get_timesteps(num_inference_steps, strength, device=device)
         latent_timestep = timesteps[:1].repeat(batch_size)
         latents = self.noise_latents_for_time(initial_latents, latent_timestep, noise_func=noise_func)
 
@@ -454,7 +454,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
 
         img2img_pipeline = StableDiffusionImg2ImgPipeline(**self.components)
         img2img_pipeline.scheduler.set_timesteps(num_inference_steps, device=device)
-        timesteps = img2img_pipeline.get_timesteps(num_inference_steps, strength, device=device)
+        timesteps, _ = img2img_pipeline.get_timesteps(num_inference_steps, strength, device=device)
 
         assert img2img_pipeline.scheduler is self.scheduler