Add support Karras sigmas for StableDiffusionKDiffusionPipeline

[takuma104]

This PR adds an option use_karras_sigmas= to apply the Karras sigmas discussed in issues like #1633, #2801 for the StableDiffusionKDiffusionPipeline. The original code is mostly by @Stax124, with some modifications made to apply it to the latest StableDiffusionKDiffusionPipeline. I opened this PR because the topic of DPM++ 2M Karras has come up repeatedly, and I personally wanted to use it.

Usage:

import torch
from diffusers import StableDiffusionKDiffusionPipeline
pipe = StableDiffusionKDiffusionPipeline.from_pretrained('dreamlike-art/dreamlike-anime-1.0', 
    torch_dtype=torch.float16,
    safety_checker = None,
).to('cuda')
pipe.set_scheduler('sample_dpmpp_2m')
image = pipe(prompt='best quality, a tuxedo cat',
    num_inference_steps=15, generator=torch.Generator(device='cuda').manual_seed(0), 
    width=512, height=512, 
    use_karras_sigmas=True).images[0]
image

Effectiveness:

The sigmas seem to change slightly like this.

Default:

[14.6146,  9.6826,  6.6780,  4.7746,  3.5221,  2.6666,  2.0606,  1.6156, 1.2768,  1.0097,  0.7913,  0.6056,  0.4397,  0.2780,  0.0292,  0.0000]

Use Karras option:

[14.6146, 10.8198,  7.9029,  5.6878,  4.0277,  2.8016,  1.9104,  1.2741, 0.8289,  0.5243,  0.3211,  0.1896,  0.1072,  0.0576,  0.0292,  0.0000]

Test:

The integration tests have passed (although there is no coverage for the additional code). For generation test, I couldn't verify pixel matching, but the results seem to be close to those of DPM++ 2M Karras in stable-diffusion-webui.

stable-diffusion-webui with DPM++ 2M Karras option	use_karras_sigmas=True	use_karras_sigmas=False

[HuggingFaceDocBuilderDev]

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

[takuma104]

For the related topic, I also looked into the question of whether the sample_dpmpp_2m option (without Karras) with StableDiffusionKDiffusionPipeline and the DPMSolverMultistepScheduler with StableDiffusionPipeline are the same.

While they appear quite similar, the results are not identical. Looking at the implementation, DPMSolverMultistepScheduler seems to be written according to the paper. In contrast, the implementation of k-diffusion did not appear to be implemented according to the formulas in the paper.

import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
pipe = StableDiffusionPipeline.from_pretrained('dreamlike-art/dreamlike-anime-1.0', 
    torch_dtype=torch.float16,
    safety_checker = None,
).to('cuda')
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
image = pipe(prompt=prompt, num_inference_steps=15, 
             generator=torch.Generator(device='cuda').manual_seed(0), 
             width=512, height=512).images[0]
image

StableDiffusionKDiffusionPipeline with sample_dpmpp_2m	StableDiffusionPipeline with DPMSolverMultistepScheduler

[patrickvonplaten]

Very clean!

Could you just run make style to merge this one @takuma104 ? :-)

[takuma104]

@patrickvonplaten thanks! I just fixed it.

[sayakpaul]

This is how docstrings should be I think.

Could you maybe just add a link to the paper that introduced Karras sigmas?

[takuma104]

@sayakpaul Thanks for merging! The so-called Karras Sigmas corresponds to Eq. 5 in the paper. Since I missed the timing, could you please include a citation as a comment when your working on #2905?

Karras, T. (2022, June 1). Elucidating the Design Space of Diffusion-Based Generative Models. arXiv.org. https://arxiv.org/abs/2206.00364

[sayakpaul]

I love everything about this PR.

Changes introduced, PR description, and the considerations given to it.

Well done!

[alexisrolland]

@takuma104 out of curiosity, why was this functionality implemented in a separate pipeline than StableDiffusionPipeline rather than as a scheduler for StableDiffusionPipeline?

[takuma104]

@alexisrolland Currently, there is a use_karras_sigmas=True option in DPMSolverMultistepScheduler (#3001), so please use it if you want to use karras_sigmas in a pipeline other than this one.

[asthinasthi]

I was looking for all the available samplers. The list is here https://github.com/crowsonkb/k-diffusion/blob/master/k_diffusion/sampling.py

sample_dpm_fast
sample_dpm_adaptive
sample_dpmpp_2s_ancestral
sample_dpmpp_sde
sample_dpmpp_2m
sample_dpmpp_2m_sde

In case anyone is interested.