[don't merge]Img2Img: timestep mismatch in scheduler with duplicated first timesteps

[yiyixuxu]

In this PR I try to explore the potential timesteps mismatch between scheduler and pipeline because how we determine the first step index in scheduler._init_step_index method. I also proposed a potential solution.

should fix #5687

the timestep mismatch issue

our schedulers use step_index to count the timestep indices, however we still need to search the index for the first timestep. In _init_step_index method, when there are duplicated timesteps, we will always take 2nd one. This will potentially cause wrong simgas to be used in scheduler.

For example, if we have 5 steps with a bunch of duplications [2, 2, ,2 ,2, 1]. If the denoising loop start from the very first step, we will send the timestep 2 to scheduler and use it to find the first sigma. The scheduler will actually start from second step because the timestep is duplicated. I think we will get a "index out of bound" error for the last step; If the denoising loop starts from the 2nd step, we will get it just right; if we start from the 3rd step, it will be able to run but there will be a mismatch.

I used below code to recreate this issue

import torch
from diffusers import DPMSolverMultistepScheduler


common_config = {'beta_start': 0.00085, 'beta_end': 0.012, 'beta_schedule': 'scaled_linear'}
scheduler =  DPMSolverMultistepScheduler(**common_config, use_karras_sigmas=True)

num_inference_steps = 100
scheduler.set_timesteps(num_inference_steps, torch.device("cpu"))

print(f"scheduler.timesteps: {scheduler.timesteps}")

for strength in [0.05, 0.04, 0.03, 0.02, 0.01 ]:
    print(" ")
    print(f" - strength: {strength} ")

    # This function implement the basic logic in img2img pipeine's `get_timesteps()` method
    # This is how we find the index for the first timestep in img2img pipeline.
    # We use this index to find the timesteps for the denoising loop
    def testing_get_timesteps(num_inference_steps, strength):
        init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
        t_start = max(num_inference_steps - init_timestep, 0)
        return t_start

    pipe_step_index_init = testing_get_timesteps(num_inference_steps, strength)
    pipe_timesteps = scheduler.timesteps[pipe_step_index_init:]
    print(f" - `timesteps` for pipeline denoising loop: {pipe_timesteps}")
    print(f" - inital step_index (pipeline): {pipe_step_index_init}")

    # this function implements the basic logic in scheduler's `_init_step_index()` method
    # this is how we find the index for first timestep in scheduler,
    # we use this index to find sigmas for denoising step in scheduler
    def testing_init_step_index(scheduler, timestep):
        index_candidates = (scheduler.timesteps == timestep).nonzero()
        if len(index_candidates) == 0:
            step_index = len(scheduler.timesteps) - 1
        elif len(index_candidates) > 1:
            step_index = index_candidates[1].item()
        else:
            step_index = index_candidates[0].item()
        return step_index

    sched_step_index_init = testing_init_step_index(scheduler,pipe_timesteps[0])
    print(f" - initial_step_index (scheduler): {sched_step_index_init}")

outputs:

scheduler.timesteps:

tensor([999, 992, 985, 978, 971, 963, 956, 948, 940, 933, 924, 916, 908, 899,
        891, 882, 873, 864, 854, 845, 835, 825, 815, 805, 794, 783, 772, 761,
        749, 737, 725, 713, 700, 687, 674, 660, 646, 632, 618, 603, 587, 572,
        556, 540, 523, 506, 489, 472, 454, 436, 418, 400, 382, 364, 345, 327,
        309, 291, 273, 256, 239, 222, 205, 190, 174, 160, 146, 133, 120, 108,
         97,  87,  78,  69,  61,  53,  47,  41,  35,  31,  26,  23,  19,  16,
         14,  12,  10,   8,   7,   6,   4,   4,   3,   2,   2,   1,   1,   1,
          0,   0])

• strength: 0.05

 - timesteps for pipeline denoising loop: tensor([1, 1, 1, 0, 0])
 - inital step_index (pipeline): 95
 - initial_step_index (scheduler): 96

• strength: 0.04

- timesteps for pipeline denoising loop: tensor([1, 1, 0, 0])
- inital step_index (pipeline): 96
- initial_step_index (scheduler): 96

• strength: 0.03

- timesteps for pipeline denoising loop: tensor([1, 0, 0])
- inital step_index (pipeline): 97
- initial_step_index (scheduler): 96

• strength: 0.02

- timesteps for pipeline denoising loop: tensor([0, 0])
- inital step_index (pipeline): 98
- initial_step_index (scheduler): 99

• strength: 0.01

- timesteps for pipeline denoising loop: tensor([0])
- inital step_index (pipeline): 99
- initial_step_index (scheduler): 99

setting initial step_index from pipeline

One solution I can think of is to set the initial step_index from the pipeline. I coded it up for stable diffusion img2img and dpm-multi scheduler as a quick example. Let me know what you think @patrickvonplaten

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[patrickvonplaten]

This needs to be a very general function call (e.g. one that every scheduler has).

We could maybe call it self.scheduler.set_begin_index(...) and then make sure that every scheduler has such a function.

[patrickvonplaten]

Let's add a setter method here as well (https://stackoverflow.com/questions/2627002/whats-the-pythonic-way-to-use-getters-and-setters) that can be set from all pipelines if necessary

[yiyixuxu]

ok but can i ask why do we need both a setter method and this set_begin_index method? they are doing exactly the same thing, no?
https://github.com/huggingface/diffusers/pull/5746/files#r1463161680

[patrickvonplaten]

Ah ok, I think we had a bit of a misunderstanding here 😅 We don't need / we shouldn't add it if it doesn't have to be used. I was under the impression that we have to use a setter method. But it seems like we don't need it after all!