[draft] refactor DPMSolverMultistepScheduler using sigmas

[yiyixuxu]

This PR refactor DPMSolverMultistepScheduler: update the computation to use sigmas.

currently, I only refactored "dpmsolver++" and "sde-dpmsolver++"
to-do left:

• "dpmsolver++"
• "sde-dpmsolver++"
• "sde-dpmsolver"
• "dpmsolver"

#4187

[HuggingFaceDocBuilderDev]

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

[yiyixuxu]

@patrickvonplaten

Do we really need "sde-dpmsolver" and "dpmsolver" here?

We have so many options here, and I think that really confuses people. i.e. we have 4 algorithm_type, 2 solver_order, and 2 solver_type, so 16 combinations here 😅

Can we possibly trim it a little bit as I refactor it? From what I understand, the algorithm type "dpmsolver" and "sde-dpmsolver" (proposed in this paper before the same author came up with dpmsolver+++ https://huggingface.co/papers/2206.00927) are completely obsolete at this point.

[yiyixuxu]

testing "dpmsolver ++"

testing the implementation against k-diffusion

• I followed the same math used in k-diffusion so results should match exact
• diff< 3e-4 for both tests below

use_karras_sigmas=False

import torch
from diffusers import StableDiffusionKDiffusionPipeline, DPMSolverMultistepScheduler, StableDiffusionPipeline
import gc
import numpy as np

# test1: use_karras_sigmas=False
seed = 33

pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
pipe = pipe.to("cuda")
pipe.scheduler = DPMSolverMultistepScheduler.from_config(
    pipe.scheduler.config, use_karras_sigmas=False
)

prompt = "an astronaut riding a horse on mars"

generator = torch.Generator(device="cuda").manual_seed(seed)
image_d = pipe(prompt, generator=generator, num_inference_steps=20, output_type='np').images[0]

seed = 33

pipe = StableDiffusionKDiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4"
).to("cuda")
pipe.set_scheduler("sample_dpmpp_2m")

prompt = "an astronaut riding a horse on mars"

generator = torch.Generator(device="cuda").manual_seed(seed)
image_k = pipe(
    prompt, generator=generator, num_inference_steps=20, use_karras_sigmas=False, output_type='np'
).images[0]

print(f"compare: {np.max(np.abs((image_d - image_k)))}")

-> 0.0002550482749938965

use_karras_sigmas=True

seed = 33

pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
pipe = pipe.to("cuda")
pipe.scheduler = DPMSolverMultistepScheduler.from_config(
    pipe.scheduler.config, use_karras_sigmas=True
)

prompt = "an astronaut riding a horse on mars"

generator = torch.Generator(device="cuda").manual_seed(seed)
image_d = pipe(prompt, generator=generator, num_inference_steps=20, output_type='np').images[0]

seed = 33

pipe = StableDiffusionKDiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4"
).to("cuda")
pipe.set_scheduler("sample_dpmpp_2m")

prompt = "an astronaut riding a horse on mars"

generator = torch.Generator(device="cuda").manual_seed(seed)
image_k = pipe(
    prompt, generator=generator, num_inference_steps=20, use_karras_sigmas=True, output_type='np'
).images[0]

print(f"compare: {np.max(np.abs((image_d - image_k)))}")

-> 0.000287860631942749

compare against current implementation (main branch)

• there are some numerical difference if we compare numpy output, but visually they are pretty identical I think

use_karras_sigmas=False

import torch
from diffusers import StableDiffusionKDiffusionPipeline, DPMSolverMultistepScheduler, StableDiffusionPipeline
import gc

seed = 33

pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
pipe = pipe.to("cuda")
pipe.scheduler = DPMSolverMultistepScheduler.from_config(
    pipe.scheduler.config, use_karras_sigmas=False
)

prompt = "an astronaut riding a horse on mars"

generator = torch.Generator(device="cuda").manual_seed(seed)
image = pipe(prompt, generator=generator, num_inference_steps=20).images[0]

image.save(f"dpmsolver++.png")

current implementation

This PR

use_karras_sigmas=True

seed = 33

pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
pipe = pipe.to("cuda")
pipe.scheduler = DPMSolverMultistepScheduler.from_config(
    pipe.scheduler.config, use_karras_sigmas=True
)

prompt = "an astronaut riding a horse on mars"

generator = torch.Generator(device="cuda").manual_seed(seed)
image = pipe(prompt, generator=generator, num_inference_steps=20).images[0]

image.save(f"dpmsolver++_k_sigma.png") 

current implementation

this PR

[yiyixuxu]

think I'm overcomplicating things with this pr - will try a different approach!