Unverified Commit 31af4d17 authored by mkshing's avatar mkshing Committed by GitHub
Browse files

Support LMSDiscreteScheduler in LDMPipeline (#891)



* Support LMSDiscreteScheduler in LDMPipeline

This is a small change to support all schedulers such as LMSDiscreteScheduler in LDMPipeline.

What's changed
-------
* Add the `scale_model_input` function before `step` to ensure correct denoising (L77)

* Add "scale the initial noise by the standard deviation required by the scheduler"

* run `make style`
Co-authored-by: default avatarAnton Lozhkov <anton@huggingface.co>
parent dec18c86
...@@ -64,6 +64,9 @@ class LDMPipeline(DiffusionPipeline): ...@@ -64,6 +64,9 @@ class LDMPipeline(DiffusionPipeline):
) )
latents = latents.to(self.device) latents = latents.to(self.device)
# scale the initial noise by the standard deviation required by the scheduler
latents = latents * self.scheduler.init_noise_sigma
self.scheduler.set_timesteps(num_inference_steps) self.scheduler.set_timesteps(num_inference_steps)
# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
...@@ -74,8 +77,9 @@ class LDMPipeline(DiffusionPipeline): ...@@ -74,8 +77,9 @@ class LDMPipeline(DiffusionPipeline):
extra_kwargs["eta"] = eta extra_kwargs["eta"] = eta
for t in self.progress_bar(self.scheduler.timesteps): for t in self.progress_bar(self.scheduler.timesteps):
latent_model_input = self.scheduler.scale_model_input(latents, t)
# predict the noise residual # predict the noise residual
noise_prediction = self.unet(latents, t).sample noise_prediction = self.unet(latent_model_input, t).sample
# compute the previous noisy sample x_t -> x_t-1 # compute the previous noisy sample x_t -> x_t-1
latents = self.scheduler.step(noise_prediction, t, latents, **extra_kwargs).prev_sample latents = self.scheduler.step(noise_prediction, t, latents, **extra_kwargs).prev_sample
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment