fix heun scheduler

[patil-suraj]

This PR fixes a subtle bug in HeunScheduler

• When computing the pred_original_sample in 2nd order, we need to use sigma_next instead of sigma_hat.

cf https://github.com/crowsonkb/k-diffusion/blob/5b3af030dd83e0297272d861c19477735d0317ec/k_diffusion/sampling.py#L180

 denoised_2 = model(x_2, sigmas[i + 1] * s_in, **extra_args)

note the sigmas[i + 1] which is passed to the model call. That's sigma_next in our implementation.

• Also when computing the derivative in 2nd order step, we need to divide by sigma_next as that's the input sigma in that step.
  cf https://github.com/crowsonkb/k-diffusion/blob/5b3af030dd83e0297272d861c19477735d0317ec/k_diffusion/sampling.py#L181

d_2 = to_d(x_2, sigmas[i + 1], denoised_2)

Not the sigmas[i + 1] passed to to_d call.

For some reason, this bug does not affect epsilon prediction mode. But generates noisy output in v_prediction with SD2 768 model.

To quickly reproduce

from diffusers import StableDiffusionPipeline, HeunDiscreteScheduler
import torch

pipe = StableDiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-2", revision="fp16", torch_dtype=torch.float16
).to("cuda")
pipe.enable_attention_slicing()

heun = HeunDiscreteScheduler.from_config(pipe.scheduler.config)
pipe.scheduler = heun

p = "A photo of an astronaut riding a horse"
torch.cuda.manual_seed(2323)
torch.manual_seed(2323)
img = pipe(p, num_inference_steps=25).images[0]
img

with main, this gives:

And with this PR:

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[patrickvonplaten]

Very cool thanks for fixing - made sure that everything is still correct for sd v1