Fix use_lu_lambdas and use_karras_sigmas with beta_schedule=squaredcos_cap_v2 in DPMSolverMultistepScheduler

[hlky]

What does this PR do?

.round() was applied with use_lu_lambdas and use_karras_sigmas, this seems to give incorrect _step_index from index_for_timestep as multiple timesteps are the same.

Reproduction

from diffusers import DPMSolverMultistepScheduler
import torch


def dummy_model():
    def model(sample, t, *args):
        # if t is a tensor, match the number of dimensions of sample
        if isinstance(t, torch.Tensor):
            num_dims = len(sample.shape)
            # pad t with 1s to match num_dims
            t = t.reshape(-1, *(1,) * (num_dims - 1)).to(
                sample.device, dtype=sample.dtype
            )

        return sample * t / (t + 1)

    return model


def dummy_sample_deter():
    batch_size = 4
    num_channels = 3
    height = 8
    width = 8

    num_elems = batch_size * num_channels * height * width
    sample = torch.arange(num_elems)
    sample = sample.reshape(num_channels, height, width, batch_size)
    sample = sample / num_elems
    sample = sample.permute(3, 0, 1, 2)

    return sample


scheduler = DPMSolverMultistepScheduler(
    **{
        "num_train_timesteps": 1000,
        "beta_start": 0.00085,
        "beta_end": 0.012,
        "beta_schedule": "squaredcos_cap_v2",
        "solver_order": 2,
        "prediction_type": "epsilon",
        "dynamic_thresholding_ratio": 0.995,
        "sample_max_value": 1.0,
        "algorithm_type": "dpmsolver++",
        "solver_type": "midpoint",
        "lower_order_final": True,
        "use_karras_sigmas": True,
        # "use_lu_lambdas": True,
        "flow_shift": 1.0,
        "final_sigmas_type": "zero",
        "timestep_spacing": "linspace",
        "steps_offset": 0,
    }
)
scheduler.set_timesteps(20)
scheduler.timesteps, scheduler.sigmas, scheduler.timesteps.shape, scheduler.sigmas.shape

model = dummy_model()
sample = dummy_sample_deter()

generator = torch.manual_seed(0)

for i, t in enumerate(scheduler.timesteps):
    print(scheduler._step_index)
    residual = model(sample, t)
    sample = scheduler.step(residual, t, sample, generator=generator).prev_sample

`use_karras_sigmas`

(tensor([999, 998, 998, 998, 998, 998, 998, 998, 998, 997, 996, 994, 989, 978,
         949, 867, 623, 221,  33,   0]),
 tensor([2.0291e+04, 1.4548e+04, 1.0258e+04, 7.1015e+03, 4.8171e+03, 3.1939e+03,
         2.0641e+03, 1.2957e+03, 7.8684e+02, 4.5987e+02, 2.5703e+02, 1.3628e+02,
         6.7820e+01, 3.1240e+01, 1.3064e+01, 4.8251e+00, 1.5101e+00, 3.7478e-01,
         6.5665e-02, 6.4271e-03, 0.0000e+00]),
 torch.Size([20]),
 torch.Size([21]))

`use_lu_lambdas`

(tensor([999, 998, 998, 998, 998, 997, 995, 991, 981, 961, 916, 820, 639, 393,
         195,  86,  35,  12,   3,   0]),
 tensor([2.0291e+04, 9.2309e+03, 4.1993e+03, 1.9103e+03, 8.6904e+02, 3.9534e+02,
         1.7985e+02, 8.1815e+01, 3.7219e+01, 1.6932e+01, 7.7024e+00, 3.5040e+00,
         1.5940e+00, 7.2514e-01, 3.2988e-01, 1.5007e-01, 6.8268e-02, 3.1056e-02,
         1.4128e-02, 6.4271e-03, 0.0000e+00]),
 torch.Size([20]),
 torch.Size([21]))

Fixes #10738

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@yiyixuxu @vladmandic

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[yiyixuxu]

I use a test like this for scheduler changes, can you run for relevant configs (need to add use_lu_lambdas?) & affected schedulers (here only DPMSolverMultistepScheduler)?

need to make sure outputs are same as in main (or better if it is expected to change)

# slow test for dpm
import torch
from diffusers import StableDiffusionXLPipeline
from diffusers import EulerDiscreteScheduler, DEISMultistepScheduler, HeunDiscreteScheduler, UniPCMultistepScheduler
import os

from diffusers.utils import make_image_grid

config_zero = {"final_sigmas_type":"zero"}

schedulers = {
    "Euler": {
        "zero": (EulerDiscreteScheduler, config_zero),
     },
     "Euler_K": {
        "zero": (EulerDiscreteScheduler, {"use_karras_sigmas": True, **config_zero}),
     },
     "Euler_EXP": {
        "zero": (EulerDiscreteScheduler, {"use_exponential_sigmas": True, **config_zero}),
     },
     "Euler_B": {
        "zero": (EulerDiscreteScheduler, {"use_beta_sigmas": True, **config_zero}),
     },
    "DEIS": {
        "zero": (DEISMultistepScheduler, config_zero),
     },
     "DEIS_K": {
        "zero": (DEISMultistepScheduler, {"use_karras_sigmas": True, **config_zero}),
     },
     "DEIS_EXP": {
        "zero": (DEISMultistepScheduler, {"use_exponential_sigmas": True, **config_zero}),
     },
     "DEIS_B": {
        "zero": (DEISMultistepScheduler, {"use_beta_sigmas": True, **config_zero}),
     },
    "Heun": {
        "zero": (HeunDiscreteScheduler, config_zero),
     },
     "Heun_K": {
        "zero": (HeunDiscreteScheduler, {"use_karras_sigmas": True, **config_zero}),
     },
     "Heun_EXP": {
        "zero": (HeunDiscreteScheduler, {"use_exponential_sigmas": True, **config_zero}),
     },
     "Heun_B": {
        "zero": (HeunDiscreteScheduler, {"use_beta_sigmas": True, **config_zero}),
     },
    "UniP": {
        "zero": (UniPCMultistepScheduler, config_zero),
     },
     "UniP_K": {
        "zero": (UniPCMultistepScheduler, {"use_karras_sigmas": True, **config_zero}),
     },
     "UniP_EXP": {
        "zero": (UniPCMultistepScheduler, {"use_exponential_sigmas": True, **config_zero}),
     },
     "UniP_B": {
        "zero": (UniPCMultistepScheduler, {"use_beta_sigmas": True, **config_zero}),
     },
}


## Test SD-XL

# model_id = "stabilityai/stable-diffusion-xl-base-1.0"
model_id = "frankjoshua/juggernautXL_version6Rundiffusion"
pipe = StableDiffusionXLPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    use_safetensors=True,
    variant="fp16",
    add_watermarker=False)
pipe = pipe.to('cuda')
prompt = "Adorable infant playing with a variety of colorful rattle toys."
save_dir = './test_juggernautxl_baby_main'

if not os.path.exists(save_dir):
    os.mkdir(save_dir)

steps = 25

params = {
    "prompt": [prompt],
    "num_inference_steps": steps,
    "guidance_scale": 7,
}
for scheduler_name in schedulers.keys():
    for seed in [123]:
        out_imgs = []
        scheduler_configs = schedulers[scheduler_name]
        for scheduler_config_name in scheduler_configs.keys():
            generator = torch.Generator(device='cuda').manual_seed(seed)
            scheduler = scheduler_configs[scheduler_config_name][0].from_pretrained(
                model_id,
                subfolder="scheduler",
                **scheduler_configs[scheduler_config_name][1],
            )
            pipe.scheduler = scheduler

            img = pipe(**params, generator=generator).images[0]
            out_imgs.append(img)
        out_img = make_image_grid(out_imgs, rows=1, cols=2)    
        out_img.save(os.path.join(save_dir, f"seed_{seed}_steps_{steps}_{scheduler_name}.png"))

[hlky]

timesteps for linear and scaled_linear would change and without round in some cases would end with 0, 0 instead of 1, 0, it's sort of the opposite issue with squaredcos_cap_v2 where round increase the first timesteps to 999, 999 instead of 999, 998 which is what causes the issue in index_for_timestep, so I've changed this PR to only round for linear and scaled_linear this keeps the existing behavior for those beta_schedule types and fixes squaredcos_cap_v2. WDYT?