Wrong amount of steps for some schedulers with base+refiner ensemble when use_karras_sigmas=True

[TimothyAlexisVass]

Describe the bug

DEISMultistepScheduler, UniPCMultistepScheduler, DPMSolverSinglestepScheduler, EulerDiscreteScheduler, DPMSolverMultistepScheduler, LMSDiscreteScheduler: Does 33 steps total instead of 40
Base: 25, Refiner: 8

HeunDiscreteScheduler: Does 33 steps total and reports strange steps
Base: 25/49 (tqdm progress), Refiner: 8

KDPM2DiscreteScheduler: Does 34 steps total and reports strange steps
Base: 26/51, Refiner: 8
also got RuntimeWarning: divide by zero encountered in log (scheduling_k_dpm_2_ancestral_discrete.py:315)

PNDMScheduler: Reports strange steps (this happens also when use_karras_sigmas=False and with any num_inference_steps it's reported as 5/6, 10/11, 15/16, etc..)
Base: 32/33, Refiner: 8

KDPM2AncestralDiscreteScheduler: Does 33 steps total and reports strange steps
Base: 25/50, Refiner: 8
also got RuntimeWarning: divide by zero encountered in log (scheduling_k_dpm_2_ancestral_discrete.py:304)

Only DDPMScheduler, EulerAncestralDiscreteScheduler and DDIMScheduler works as expected out of the ones I tested...
Meaning, they all did 32 steps with the base model and 8 steps with the refiner.

Have gotten divide by zero warning in more than those reported here... So I added a PR #5543 for that.

Reproduction

from diffusers import AutoencoderKL, StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
import torch

vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
base = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0",
    vae=vae,
    torch_dtype=torch.float16,
    variant="fp16",
    use_safetensors=True,
)
_ = base.to("cuda")

refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0",
    vae=vae,
    text_encoder_2=base.text_encoder_2,
    torch_dtype=torch.float16,
    variant="fp16",
    use_safetensors=True,
)

_ = refiner.to("cuda")

prompt="LOVE"
num_inference_steps = 40

for scheduler in base.scheduler.compatibles:
    scheduler_name = scheduler.__name__
    if scheduler_name not in ("DPMSolverSDEScheduler"):
        refiner.scheduler = base.scheduler = scheduler.from_config(base.scheduler.config, use_karras_sigmas=True)
        print("Generating with", scheduler_name)
        latents = base(prompt, num_inference_steps=num_inference_steps, denoising_end=0.8, output_type="latent").images
        image = refiner(prompt, num_inference_steps=num_inference_steps, denoising_start=0.8, image=latents).images[0]
        display(image)

Logs

No response

System Info

diffusers 0.22.0.dev0
python 3.10.12

Who can help?

@yiyixuxu @DN6 @patrickvonplaten @sayakpaul @patrickvonplaten

[DN6]

Hi @TimothyAlexisVass could you provide a reproducible code example to demonstrate the issue please?

[TimothyAlexisVass]

could you provide a reproducible code example
Yes, I've added it above 😄 👍

[DN6]

@TimothyAlexisVass I was able to reproduce this.

cc: @yiyixuxu Could you please take a look here.

[TimothyAlexisVass]

This is how I'm currently working around the issue... By multiplying the step value by 1.2 for certain schedulers.
The value within parenthesis is the actual steps wanted and the one without parenthesis is the multiplied value that is sent as num_inference_steps.

[patrickvonplaten]

I'm wondering whether the problem is just the progress bar here actually

[yiyixuxu]

hey @TimothyAlexisVass
For base+refiner ensemble, we need to use same scheduler for both base and refiner in order for it to work correctly

cc @patrickvonplaten, can you confirm that it is not a bug in our implementation, i.e. we don't actually need to consider the use case where you use different schedulers for base and refiner

[TimothyAlexisVass]

The reason why I use different schedulers was because for some schedulers, it's not possible to get any result by using the same scheduler on both base and refiner...

When making this change to the example:
refiner.scheduler = base.scheduler = scheduler.from_config(base.scheduler.config, use_karras_sigmas=True)
Then both base and refiner use the same scheduler.
(I've updated the reproducible example above too...)

Notice here for example that the progress bars are wrong and the image is black.
There is a similar result with KDPM2DiscreteScheduler... Which will generate an image, but it doesn't get denoised.

---

Here is another example, even though it generates an image, it does the switch to refiner at 0.625 rather than 0.8
This also happens for DMPSolverSinglestepScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler.

---

HeunScheduler exhibits a combination of that... wrong steps reported, switches to 0.625 and generates a black image.

---

This one doesn't get past step 2...

RuntimeError: cusolver error: CUSOLVER_STATUS_INTERNAL_ERROR, when calling cusolverDnCreate(handle)

---

PNDMScheduler does the same thing with reporting progress 32/33 for base.

Conclusion:
Still, only DDPMScheduler, EulerAncestralDiscreteScheduler and DDIMScheduler works as expected...
Meaning, they all did 32 steps with the base model and 8 steps with the refiner while reporting progress correctly.

Sidenote:
Perhaps it's possible to implement karras sigmas for DEISMultistepScheduler too?
AttributeError: 'DEISMultistepScheduler' object has no attribute '_convert_to_karras'

[yiyixuxu]

@TimothyAlexisVass

1. you should use same scheduler for both base and refiner
2. the 0.8 ratio applies to training steps - meaning that 80% of the training steps belong to the high-noise stage and 20% belong to the low noise stage. Depends on how your scheduler space the inference steps, we may not see the same ratio in the inference steps between base and refiner. I think the 0.625 split you observed is expected
3. so I think only remaining the issue is it generates black image with some schedulers. I'm happy to look into this! If you can update the PR with reproducible script for schedulers that produce black images that will be great:)

[TimothyAlexisVass]

Hi @yiyixuxu

1. The second test run with the images included was with the same scheduler for both base and refiner. So, the reported results are with the same scheduler for both models. I had also updated the reproducible example code for the second example, so you can use that.
2. Yes, it's true that some schedulers have a spacing which sets timestep 200 around 0.625 instead of 0.8. Thanks for clarifying that.
3. I think that every issue reported needs to be fixed, which are:
   3a. Some schedulers generating black images.
   3b. KDPM2DiscreteScheduler generating an unfinished result. (See image below)
   3c. UniPC crashing at step 2.
   3d. Wrong amount of (total) steps reported in the progress bar.

[yiyixuxu]

@TimothyAlexisVass
are you using the update-to-date code? I think we merged in a PR very recently that should have fixed the issue with second order schedulers #5526

let me know:)

[TimothyAlexisVass]

I have done like this:

pip install --upgrade -qq git+https://github.com/huggingface/diffusers

With version 0.23.0.dev0:

• UniPCMultistepScheduler works correctly.
• HeunDiscreteScheduler, KDPM2AncestralDiscrete and KDPM2DiscreteScheduler generate images.

So, only the total step count needs to be fixed for the second order schedulers, as well as PNDMScheduler.

Seems they sometimes do an extra step too (11 total, num_inference_steps was 10).

Still reports 1 more than what it actually does for the base model.