New Scheduler: add Euler Ancestral Scheduler to StableDiffusionPipeline

src/diffusers/__init__.py

@@ -37,6 +37,7 @@
         PNDMScheduler,
         SchedulerMixin,
         ScoreSdeVeScheduler,
+        EulerAScheduler,
     )
     from .training_utils import EMAModel
 else:

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py

@@ -8,7 +8,7 @@
 from ...configuration_utils import FrozenDict
 from ...models import AutoencoderKL, UNet2DConditionModel
 from ...pipeline_utils import DiffusionPipeline
-from ...schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler
+from ...schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, EulerAScheduler
 from ...utils import deprecate, logging
 from . import StableDiffusionPipelineOutput
 from .safety_checker import StableDiffusionSafetyChecker
@@ -306,22 +306,43 @@ def __call__(
         extra_step_kwargs = {}
         if accepts_eta:
             extra_step_kwargs["eta"] = eta
+        if generator is not None:
+            extra_step_kwargs["generator"] = generator
 
         for i, t in enumerate(self.progress_bar(timesteps_tensor)):
             # expand the latents if we are doing classifier free guidance
             latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
-            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+            # latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+            #TODO merge conform the EulerAScheduler interface to the standardized scheduler interface 
+            if isinstance(self.scheduler, EulerAScheduler): 
+                latent_unscaled = latent_model_input # store the unscaled latent
+                c_out, c_in, sigma_in = self.scheduler.prepare_input(latent_model_input, t)
+                # latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+                latent_model_input = latent_unscaled * c_in
+                # sigma_in = self.scheduler.get_sigma_in(latent_model_input,t)
+                eps = self.unet(latent_model_input, sigma_in , encoder_hidden_states=text_embeddings).sample
+                noise_pred = latent_unscaled + eps * c_out
 
-            # predict the noise residual
-            noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
+            else:
+                # predict the noise residual
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
 
             # perform guidance
             if do_classifier_free_guidance:
                 noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                 noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
 
+            if isinstance(self.scheduler, EulerAScheduler):
+                # change from self.scheduler.timesteps.shape[0] - 1 to num_inference_steps
+                if i <  self.scheduler.num_inference_steps: #avoid out of bound error
+                    # t_prev = self.scheduler.timesteps[i+1]
+                    latents = self.scheduler.step(noise_pred, i, latents, **extra_step_kwargs).prev_sample
+            else: 
             # compute the previous noisy sample x_t -> x_t-1
-            latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
+                    latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
 
             # call the callback, if provided
             if callback is not None and i % callback_steps == 0:

src/diffusers/schedulers/__init__.py

@@ -24,6 +24,7 @@
     from .scheduling_sde_ve import ScoreSdeVeScheduler
     from .scheduling_sde_vp import ScoreSdeVpScheduler
     from .scheduling_utils import SchedulerMixin
+    from .scheduling_euler_a import EulerAScheduler
 else:
     from ..utils.dummy_pt_objects import *  # noqa F403