1 files changed, 5 insertions, 30 deletions

diff --git a/pipelines/stable_diffusion/vlpn_stable_diffusion.py b/pipelines/stable_diffusion/vlpn_stable_diffusion.py
index 8fbe5f9..a198cf6 100644
--- a/pipelines/stable_diffusion/vlpn_stable_diffusion.py
+++ b/pipelines/stable_diffusion/vlpn_stable_diffusion.py
@@ -216,7 +216,6 @@ class VlpnStableDiffusion(DiffusionPipeline):
 
         offset = self.scheduler.config.get("steps_offset", 0)
         init_timestep = num_inference_steps + offset
-        ensure_sigma = not isinstance(latents, PIL.Image.Image)
 
         # get the initial random noise unless the user supplied it
 
@@ -246,13 +245,8 @@ class VlpnStableDiffusion(DiffusionPipeline):
             init_timestep = int(num_inference_steps * strength) + offset
             init_timestep = min(init_timestep, num_inference_steps)
 
-            if isinstance(self.scheduler, LMSDiscreteScheduler):
-                timesteps = torch.tensor(
-                    [num_inference_steps - init_timestep] * batch_size, dtype=torch.long, device=self.device
-                )
-            else:
-                timesteps = self.scheduler.timesteps[-init_timestep]
-                timesteps = torch.tensor([timesteps] * batch_size, dtype=torch.long, device=self.device)
+            timesteps = self.scheduler.timesteps[-init_timestep]
+            timesteps = torch.tensor([timesteps] * batch_size, device=self.device)
 
             # add noise to latents using the timesteps
             noise = torch.randn(latents.shape, generator=generator, device=self.device)
@@ -263,13 +257,6 @@ class VlpnStableDiffusion(DiffusionPipeline):
             if latents.device != self.device:
                 raise ValueError(f"Unexpected latents device, got {latents.device}, expected {self.device}")
 
-        # if we use LMSDiscreteScheduler, let's make sure latents are multiplied by sigmas
-        if ensure_sigma:
-            if isinstance(self.scheduler, LMSDiscreteScheduler):
-                latents = latents * self.scheduler.sigmas[0]
-            elif isinstance(self.scheduler, EulerAScheduler):
-                latents = latents * self.scheduler.sigmas[0]
-
         t_start = max(num_inference_steps - init_timestep + offset, 0)
 
         # Some schedulers like PNDM have timesteps as arrays
@@ -290,19 +277,13 @@ class VlpnStableDiffusion(DiffusionPipeline):
             extra_step_kwargs["generator"] = generator
 
         for i, t in enumerate(self.progress_bar(timesteps_tensor)):
-            t_index = t_start + i
-
             # expand the latents if we are doing classifier free guidance
             latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
-
-            if isinstance(self.scheduler, LMSDiscreteScheduler):
-                sigma = self.scheduler.sigmas[t_index]
-                # the model input needs to be scaled to match the continuous ODE formulation in K-LMS
-                latent_model_input = latent_model_input / ((sigma**2 + 1) ** 0.5)
+            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
 
             noise_pred = None
             if isinstance(self.scheduler, EulerAScheduler):
-                sigma = self.scheduler.sigmas[t].reshape(1)
+                sigma = t.reshape(1)
                 sigma_in = torch.cat([sigma] * latent_model_input.shape[0])
                 noise_pred = CFGDenoiserForward(self.unet, latent_model_input, sigma_in,
                                                 text_embeddings, guidance_scale, quantize=True, DSsigmas=self.scheduler.DSsigmas)
@@ -316,13 +297,7 @@ class VlpnStableDiffusion(DiffusionPipeline):
                 noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
 
             # compute the previous noisy sample x_t -> x_t-1
-            if isinstance(self.scheduler, LMSDiscreteScheduler):
-                latents = self.scheduler.step(noise_pred, t_index, latents, **extra_step_kwargs).prev_sample
-            elif isinstance(self.scheduler, EulerAScheduler):
-                latents = self.scheduler.step(noise_pred, t_index, t_index + 1,
-                                              latents, **extra_step_kwargs).prev_sample
-            else:
-                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
 
         # scale and decode the image latents with vae
         latents = 1 / 0.18215 * latents