Update LCMScheduler Inference Timesteps to be More Evenly Spaced

[dg845]

What does this PR do?

This PR modifies LCMScheduler.set_timesteps to produce timesteps which are more evenly spaced across the range of candidate timesteps, which should lead to better sample quality.

Fixes #5815.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

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.

@patrickvonplaten
@patil-suraj
@aifartist
@luosiallen

[dg845]

The solution I have currently implemented differs from the solution suggested in #5815 (comment) as follows:

• Instead of generating timesteps directly, we select indices from the training/distillation schedule lcm_origin_timesteps. This should ensure that we always end up with timesteps that we saw during training/distillation.
• Because np.linspace(0, len(lcm_origin_timesteps) - 1, num=num_inference_steps) is used to generate the indices, the first and last timesteps in the final inference timestep schedule are always anchored to the first and last timesteps in lcm_origin_timesteps. (Under default settings, this would be 999 and 19 respectively.)

[dg845]

(Also, note that the timestep schedule produced by LCMScheduler.set_timesteps will in general be different after the change, so any test that uses LCMScheduler with num_inference_steps > 1 will probably need to be fixed.)

[luosiallen]

I think we want to avoid the cases that 3-step sampling is very similar to 2-step sampling, which means that we don't want setting anchor for the first timesteps 19. For example, i think the schedule of 2-step [999, 499] is better than [999, 19].

[patil-suraj]

Thanks a lot for fixing this @dg845 , lgtm!

Will run some tests then we can merge.

[HuggingFaceDocBuilderDev]

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

[patrickvonplaten]

I trust you here @patil-suraj

[patrickvonplaten]

If you consider the previous timestep spacing a bug, I'm ok with the PR. If not and it's just a "different way" of doing timestep spacing, we should probably be careful about backwards breaking

[luosiallen]

Should we do more test? or just like @patrickvonplaten suggestion. We can just add timestep [list] input in the __call__, so the users can figure out what's the best for them.

[patil-suraj]

I'm currently running the slow tests that we have. All the fast tests are passing.
This timestep schedule is correct, as it matches the schedule during training, so it should be the default IMO. We could add the timestep argument to __call__ for additional schedules.

[luosiallen]

I'm currently running the slow tests that we have. All the fast tests are passing.
This timestep schedule is correct, as it matches the schedule during training, so it should be the default IMO. We could add the timestep argument to __call__ for additional schedules.

Thanks @patil-suraj !

[patil-suraj]

Some slow tests are failing,

• pytest tests/lora/test_lora_layers_peft.py::LoraSDXLIntegrationTests::test_sdxl_lcm_lora
• pytest tests/lora/test_lora_layers_peft.py::LoraSDXLIntegrationTests::test_sdv1_5_lcm_lora
• pytest tests/lora/test_lora_layers_peft.py::LoraSDXLIntegrationTests::test_sdv1_5_lcm_lora_img2img

@patrickvonplaten would it be okay to adjust the tests in this case ?

[patil-suraj]

All slow LCM tests are now passing, I updated the images here https://huggingface.co/datasets/hf-internal-testing/diffusers-images/tree/main/lcm_lora